Skip to content

Commit 30b1ae5

Browse files
author
Frank
committed
zen: rate limit
1 parent 9cd465f commit 30b1ae5

29 files changed

Lines changed: 692 additions & 491 deletions

File tree

‎bun.lock‎

Lines changed: 36 additions & 17 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎github/sst-env.d.ts‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -6,4 +6,4 @@
66
/// <reference path="../sst-env.d.ts" />
77

88
import "sst"
9-
export {}
9+
export {}

‎infra/console.ts‎

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -106,6 +106,7 @@ const AUTH_API_URL = new sst.Linkable("AUTH_API_URL", {
106106
const STRIPE_WEBHOOK_SECRET = new sst.Linkable("STRIPE_WEBHOOK_SECRET", {
107107
properties: { value: stripeWebhook.secret },
108108
})
109+
const gatewayKv = new sst.cloudflare.Kv("GatewayKv")
109110

110111
////////////////
111112
// CONSOLE
@@ -136,6 +137,16 @@ new sst.cloudflare.x.SolidStart("Console", {
136137
EMAILOCTOPUS_API_KEY,
137138
AWS_SES_ACCESS_KEY_ID,
138139
AWS_SES_SECRET_ACCESS_KEY,
140+
...($dev
141+
? [
142+
new sst.Secret(
143+
"CLOUDFLARE_DEFAULT_ACCOUNT_ID",
144+
process.env.CLOUDFLARE_DEFAULT_ACCOUNT_ID!,
145+
),
146+
new sst.Secret("CLOUDFLARE_API_TOKEN", process.env.CLOUDFLARE_API_TOKEN!),
147+
]
148+
: []),
149+
gatewayKv,
139150
],
140151
environment: {
141152
//VITE_DOCS_URL: web.url.apply((url) => url!),

‎package.json‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -52,7 +52,7 @@
5252
"@tsconfig/bun": "catalog:",
5353
"husky": "9.1.7",
5454
"prettier": "3.6.2",
55-
"sst": "3.17.22",
55+
"sst": "3.17.23",
5656
"turbo": "2.5.6"
5757
},
5858
"dependencies": {

‎packages/console/app/src/routes/zen/util/error.ts‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,3 +3,4 @@ export class CreditsError extends Error {}
33
export class MonthlyLimitError extends Error {}
44
export class UserLimitError extends Error {}
55
export class ModelError extends Error {}
6+
export class RateLimitError extends Error {}

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

Lines changed: 40 additions & 33 deletions
Original file line numberDiff line numberDiff line change
@@ -12,14 +12,14 @@ import { UserTable } from "@opencode-ai/console-core/schema/user.sql.js"
1212
import { ModelTable } from "@opencode-ai/console-core/schema/model.sql.js"
1313
import { ProviderTable } from "@opencode-ai/console-core/schema/provider.sql.js"
1414
import { logger } from "./logger"
15-
import { AuthError, CreditsError, MonthlyLimitError, UserLimitError, ModelError } from "./error"
15+
import { AuthError, CreditsError, MonthlyLimitError, UserLimitError, ModelError, RateLimitError } from "./error"
1616
import { createBodyConverter, createStreamPartConverter, createResponseConverter } from "./provider/provider"
1717
import { anthropicHelper } from "./provider/anthropic"
1818
import { openaiHelper } from "./provider/openai"
1919
import { oaCompatHelper } from "./provider/openai-compatible"
20+
import { createRateLimiter } from "./rateLimiter"
2021

2122
type ZenData = Awaited<ReturnType<typeof ZenData.list>>
22-
type Model = ZenData["models"][string]
2323

2424
export async function handler(
2525
input: APIEvent,
@@ -28,23 +28,30 @@ export async function handler(
2828
parseApiKey: (headers: Headers) => string | undefined
2929
},
3030
) {
31+
type AuthInfo = Awaited<ReturnType<typeof authenticate>>
32+
type ModelInfo = Awaited<ReturnType<typeof validateModel>>
33+
type ProviderInfo = Awaited<ReturnType<typeof selectProvider>>
34+
3135
const FREE_WORKSPACES = [
3236
"wrk_01K46JDFR0E75SG2Q8K172KF3Y", // frank
3337
"wrk_01K6W1A3VE0KMNVSCQT43BG2SX", // opencode bench
3438
]
3539

3640
try {
3741
const body = await input.request.json()
42+
const ip = input.request.headers.get("x-real-ip") ?? ""
3843
logger.metric({
3944
is_tream: !!body.stream,
4045
session: input.request.headers.get("x-opencode-session"),
4146
request: input.request.headers.get("x-opencode-request"),
4247
})
4348
const zenData = ZenData.list()
4449
const modelInfo = validateModel(zenData, body.model)
45-
const providerInfo = selectProvider(zenData, modelInfo, input.request.headers.get("x-real-ip") ?? "")
50+
const providerInfo = selectProvider(zenData, modelInfo, ip)
4651
const authInfo = await authenticate(modelInfo, providerInfo)
47-
validateBilling(modelInfo, authInfo)
52+
const rateLimiter = createRateLimiter(modelInfo.id, modelInfo.rateLimit, ip)
53+
await rateLimiter?.check()
54+
validateBilling(authInfo, modelInfo)
4855
validateModelSettings(authInfo)
4956
updateProviderKey(authInfo, providerInfo)
5057
logger.metric({ provider: providerInfo.id })
@@ -59,7 +66,7 @@ export async function handler(
5966
}),
6067
)
6168
logger.debug("REQUEST URL: " + reqUrl)
62-
logger.debug("REQUEST: " + reqBody)
69+
logger.debug("REQUEST: " + reqBody.substring(0, 300) + "...")
6370
const res = await fetch(reqUrl, {
6471
method: "POST",
6572
headers: (() => {
@@ -84,9 +91,6 @@ export async function handler(
8491
}
8592
}
8693
logger.debug("STATUS: " + res.status + " " + res.statusText)
87-
if (res.status === 400 || res.status === 503) {
88-
logger.debug("RESPONSE: " + (await res.text()))
89-
}
9094

9195
// Handle non-streaming response
9296
if (!body.stream) {
@@ -95,6 +99,7 @@ export async function handler(
9599
const body = JSON.stringify(responseConverter(json))
96100
logger.metric({ response_length: body.length })
97101
logger.debug("RESPONSE: " + body)
102+
await rateLimiter?.track()
98103
await trackUsage(authInfo, modelInfo, providerInfo, json.usage)
99104
await reload(authInfo)
100105
return new Response(body, {
@@ -123,6 +128,7 @@ export async function handler(
123128
response_length: responseLength,
124129
"timestamp.last_byte": Date.now(),
125130
})
131+
await rateLimiter?.track()
126132
const usage = usageParser.retrieve()
127133
if (usage) {
128134
await trackUsage(authInfo, modelInfo, providerInfo, usage)
@@ -197,6 +203,15 @@ export async function handler(
197203
{ status: 401 },
198204
)
199205

206+
if (error instanceof RateLimitError)
207+
return new Response(
208+
JSON.stringify({
209+
type: "error",
210+
error: { type: error.constructor.name, message: error.message },
211+
}),
212+
{ status: 429 },
213+
)
214+
200215
return new Response(
201216
JSON.stringify({
202217
type: "error",
@@ -221,8 +236,8 @@ export async function handler(
221236
return { id: modelId, ...modelData }
222237
}
223238

224-
function selectProvider(zenData: ZenData, model: Awaited<ReturnType<typeof validateModel>>, ip: string) {
225-
const providers = model.providers
239+
function selectProvider(zenData: ZenData, modelInfo: ModelInfo, ip: string) {
240+
const providers = modelInfo.providers
226241
.filter((provider) => !provider.disabled)
227242
.flatMap((provider) => Array<typeof provider>(provider.weight ?? 1).fill(provider))
228243

@@ -235,22 +250,22 @@ export async function handler(
235250
throw new ModelError(`Provider ${provider.id} not supported`)
236251
}
237252

238-
const format = zenData.providers[provider.id].format
239-
240253
return {
241254
...provider,
242255
...zenData.providers[provider.id],
243-
...(format === "anthropic" ? anthropicHelper : format === "openai" ? openaiHelper : oaCompatHelper),
256+
...(() => {
257+
const format = zenData.providers[provider.id].format
258+
if (format === "anthropic") return anthropicHelper
259+
if (format === "openai") return openaiHelper
260+
return oaCompatHelper
261+
})(),
244262
}
245263
}
246264

247-
async function authenticate(
248-
model: Awaited<ReturnType<typeof validateModel>>,
249-
providerInfo: Awaited<ReturnType<typeof selectProvider>>,
250-
) {
265+
async function authenticate(modelInfo: ModelInfo, providerInfo: ProviderInfo) {
251266
const apiKey = opts.parseApiKey(input.request.headers)
252267
if (!apiKey) {
253-
if (model.allowAnonymous) return
268+
if (modelInfo.allowAnonymous) return
254269
throw new AuthError("Missing API key.")
255270
}
256271

@@ -282,7 +297,7 @@ export async function handler(
282297
.innerJoin(WorkspaceTable, eq(WorkspaceTable.id, KeyTable.workspaceID))
283298
.innerJoin(BillingTable, eq(BillingTable.workspaceID, KeyTable.workspaceID))
284299
.innerJoin(UserTable, and(eq(UserTable.workspaceID, KeyTable.workspaceID), eq(UserTable.id, KeyTable.userID)))
285-
.leftJoin(ModelTable, and(eq(ModelTable.workspaceID, KeyTable.workspaceID), eq(ModelTable.model, model.id)))
300+
.leftJoin(ModelTable, and(eq(ModelTable.workspaceID, KeyTable.workspaceID), eq(ModelTable.model, modelInfo.id)))
286301
.leftJoin(
287302
ProviderTable,
288303
and(eq(ProviderTable.workspaceID, KeyTable.workspaceID), eq(ProviderTable.provider, providerInfo.id)),
@@ -308,11 +323,11 @@ export async function handler(
308323
}
309324
}
310325

311-
function validateBilling(model: Model, authInfo: Awaited<ReturnType<typeof authenticate>>) {
326+
function validateBilling(authInfo: AuthInfo, modelInfo: ModelInfo) {
312327
if (!authInfo) return
313328
if (authInfo.provider?.credentials) return
314329
if (authInfo.isFree) return
315-
if (model.allowAnonymous) return
330+
if (modelInfo.allowAnonymous) return
316331

317332
const billing = authInfo.billing
318333
if (!billing.paymentMethodID)
@@ -356,26 +371,18 @@ export async function handler(
356371
}
357372
}
358373

359-
function validateModelSettings(authInfo: Awaited<ReturnType<typeof authenticate>>) {
374+
function validateModelSettings(authInfo: AuthInfo) {
360375
if (!authInfo) return
361376
if (authInfo.isDisabled) throw new ModelError("Model is disabled")
362377
}
363378

364-
function updateProviderKey(
365-
authInfo: Awaited<ReturnType<typeof authenticate>>,
366-
providerInfo: Awaited<ReturnType<typeof selectProvider>>,
367-
) {
379+
function updateProviderKey(authInfo: AuthInfo, providerInfo: ProviderInfo) {
368380
if (!authInfo) return
369381
if (!authInfo.provider?.credentials) return
370382
providerInfo.apiKey = authInfo.provider.credentials
371383
}
372384

373-
async function trackUsage(
374-
authInfo: Awaited<ReturnType<typeof authenticate>>,
375-
modelInfo: ReturnType<typeof validateModel>,
376-
providerInfo: Awaited<ReturnType<typeof selectProvider>>,
377-
usage: any,
378-
) {
385+
async function trackUsage(authInfo: AuthInfo, modelInfo: ModelInfo, providerInfo: ProviderInfo, usage: any) {
379386
const { inputTokens, outputTokens, reasoningTokens, cacheReadTokens, cacheWrite5mTokens, cacheWrite1hTokens } =
380387
providerInfo.normalizeUsage(usage)
381388

@@ -483,7 +490,7 @@ export async function handler(
483490
)
484491
}
485492

486-
async function reload(authInfo: Awaited<ReturnType<typeof authenticate>>) {
493+
async function reload(authInfo: AuthInfo) {
487494
if (!authInfo) return
488495
if (authInfo.isFree) return
489496
if (authInfo.provider?.credentials) return

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

Lines changed: 36 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,36 @@
1+
import { Resource } from "@opencode-ai/console-resource"
2+
import { RateLimitError } from "./error"
3+
import { logger } from "./logger"
4+
5+
export function createRateLimiter(model: string, limit: number | undefined, ip: string) {
6+
if (!limit) return
7+
8+
const now = Date.now()
9+
const currKey = `usage:${ip}:${model}:${buildYYYYMMDDHH(now)}`
10+
const prevKey = `usage:${ip}:${model}:${buildYYYYMMDDHH(now - 3_600_000)}`
11+
let currRate: number
12+
let prevRate: number
13+
14+
return {
15+
track: async () => {
16+
await Resource.GatewayKv.put(currKey, currRate + 1, { expirationTtl: 3600 })
17+
},
18+
check: async () => {
19+
const values = await Resource.GatewayKv.get([currKey, prevKey])
20+
const prevValue = values?.get(prevKey)
21+
const currValue = values?.get(currKey)
22+
prevRate = prevValue ? parseInt(prevValue) : 0
23+
currRate = currValue ? parseInt(currValue) : 0
24+
logger.debug(`rate limit ${model} prev/curr: ${prevRate}/${currRate}`)
25+
if (prevRate + currRate >= limit)
26+
throw new RateLimitError(`Rate limit exceeded. Please try again later.`)
27+
},
28+
}
29+
}
30+
31+
function buildYYYYMMDDHH(timestamp: number) {
32+
return new Date(timestamp)
33+
.toISOString()
34+
.replace(/[^0-9]/g, "")
35+
.substring(0, 10)
36+
}

‎packages/console/app/sst-env.d.ts‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -6,4 +6,4 @@
66
/// <reference path="../../../sst-env.d.ts" />
77

88
import "sst"
9-
export {}
9+
export {}

‎packages/console/core/script/update-models.ts‎

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,6 @@ import { ZenData } from "../src/model"
77

88
const root = path.resolve(process.cwd(), "..", "..", "..")
99
const models = await $`bun sst secret list`.cwd(root).text()
10-
console.log("models", models)
1110

1211
// read the line starting with "ZEN_MODELS"
1312
const oldValue1 = models

‎packages/console/core/src/model.ts‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,7 @@ export namespace ZenData {
2424
cost: ModelCostSchema,
2525
cost200K: ModelCostSchema.optional(),
2626
allowAnonymous: z.boolean().optional(),
27+
rateLimit: z.number().optional(),
2728
providers: z.array(
2829
z.object({
2930
id: z.string(),

0 commit comments

Comments
 (0)