diff options
Diffstat (limited to 'packages/console/app/src')
| -rw-r--r-- | packages/console/app/src/routes/zen/util/error.ts | 1 | ||||
| -rw-r--r-- | packages/console/app/src/routes/zen/util/handler.ts | 73 | ||||
| -rw-r--r-- | packages/console/app/src/routes/zen/util/rateLimiter.ts | 36 |
3 files changed, 77 insertions, 33 deletions
diff --git a/packages/console/app/src/routes/zen/util/error.ts b/packages/console/app/src/routes/zen/util/error.ts index dfc7e9fcd..f1e131468 100644 --- a/packages/console/app/src/routes/zen/util/error.ts +++ b/packages/console/app/src/routes/zen/util/error.ts @@ -3,3 +3,4 @@ export class CreditsError extends Error {} export class MonthlyLimitError extends Error {} export class UserLimitError extends Error {} export class ModelError extends Error {} +export class RateLimitError extends Error {} diff --git a/packages/console/app/src/routes/zen/util/handler.ts b/packages/console/app/src/routes/zen/util/handler.ts index 194a7c71e..edaac3a7b 100644 --- a/packages/console/app/src/routes/zen/util/handler.ts +++ b/packages/console/app/src/routes/zen/util/handler.ts @@ -12,14 +12,14 @@ import { UserTable } from "@opencode-ai/console-core/schema/user.sql.js" import { ModelTable } from "@opencode-ai/console-core/schema/model.sql.js" import { ProviderTable } from "@opencode-ai/console-core/schema/provider.sql.js" import { logger } from "./logger" -import { AuthError, CreditsError, MonthlyLimitError, UserLimitError, ModelError } from "./error" +import { AuthError, CreditsError, MonthlyLimitError, UserLimitError, ModelError, RateLimitError } from "./error" import { createBodyConverter, createStreamPartConverter, createResponseConverter } from "./provider/provider" import { anthropicHelper } from "./provider/anthropic" import { openaiHelper } from "./provider/openai" import { oaCompatHelper } from "./provider/openai-compatible" +import { createRateLimiter } from "./rateLimiter" type ZenData = Awaited<ReturnType<typeof ZenData.list>> -type Model = ZenData["models"][string] export async function handler( input: APIEvent, @@ -28,6 +28,10 @@ export async function handler( parseApiKey: (headers: Headers) => string | undefined }, ) { + type AuthInfo = Awaited<ReturnType<typeof authenticate>> + type ModelInfo = Awaited<ReturnType<typeof validateModel>> + type ProviderInfo = Awaited<ReturnType<typeof selectProvider>> + const FREE_WORKSPACES = [ "wrk_01K46JDFR0E75SG2Q8K172KF3Y", // frank "wrk_01K6W1A3VE0KMNVSCQT43BG2SX", // opencode bench @@ -35,6 +39,7 @@ export async function handler( try { const body = await input.request.json() + const ip = input.request.headers.get("x-real-ip") ?? "" logger.metric({ is_tream: !!body.stream, session: input.request.headers.get("x-opencode-session"), @@ -42,9 +47,11 @@ export async function handler( }) const zenData = ZenData.list() const modelInfo = validateModel(zenData, body.model) - const providerInfo = selectProvider(zenData, modelInfo, input.request.headers.get("x-real-ip") ?? "") + const providerInfo = selectProvider(zenData, modelInfo, ip) const authInfo = await authenticate(modelInfo, providerInfo) - validateBilling(modelInfo, authInfo) + const rateLimiter = createRateLimiter(modelInfo.id, modelInfo.rateLimit, ip) + await rateLimiter?.check() + validateBilling(authInfo, modelInfo) validateModelSettings(authInfo) updateProviderKey(authInfo, providerInfo) logger.metric({ provider: providerInfo.id }) @@ -59,7 +66,7 @@ export async function handler( }), ) logger.debug("REQUEST URL: " + reqUrl) - logger.debug("REQUEST: " + reqBody) + logger.debug("REQUEST: " + reqBody.substring(0, 300) + "...") const res = await fetch(reqUrl, { method: "POST", headers: (() => { @@ -84,9 +91,6 @@ export async function handler( } } logger.debug("STATUS: " + res.status + " " + res.statusText) - if (res.status === 400 || res.status === 503) { - logger.debug("RESPONSE: " + (await res.text())) - } // Handle non-streaming response if (!body.stream) { @@ -95,6 +99,7 @@ export async function handler( const body = JSON.stringify(responseConverter(json)) logger.metric({ response_length: body.length }) logger.debug("RESPONSE: " + body) + await rateLimiter?.track() await trackUsage(authInfo, modelInfo, providerInfo, json.usage) await reload(authInfo) return new Response(body, { @@ -123,6 +128,7 @@ export async function handler( response_length: responseLength, "timestamp.last_byte": Date.now(), }) + await rateLimiter?.track() const usage = usageParser.retrieve() if (usage) { await trackUsage(authInfo, modelInfo, providerInfo, usage) @@ -197,6 +203,15 @@ export async function handler( { status: 401 }, ) + if (error instanceof RateLimitError) + return new Response( + JSON.stringify({ + type: "error", + error: { type: error.constructor.name, message: error.message }, + }), + { status: 429 }, + ) + return new Response( JSON.stringify({ type: "error", @@ -221,8 +236,8 @@ export async function handler( return { id: modelId, ...modelData } } - function selectProvider(zenData: ZenData, model: Awaited<ReturnType<typeof validateModel>>, ip: string) { - const providers = model.providers + function selectProvider(zenData: ZenData, modelInfo: ModelInfo, ip: string) { + const providers = modelInfo.providers .filter((provider) => !provider.disabled) .flatMap((provider) => Array<typeof provider>(provider.weight ?? 1).fill(provider)) @@ -235,22 +250,22 @@ export async function handler( throw new ModelError(`Provider ${provider.id} not supported`) } - const format = zenData.providers[provider.id].format - return { ...provider, ...zenData.providers[provider.id], - ...(format === "anthropic" ? anthropicHelper : format === "openai" ? openaiHelper : oaCompatHelper), + ...(() => { + const format = zenData.providers[provider.id].format + if (format === "anthropic") return anthropicHelper + if (format === "openai") return openaiHelper + return oaCompatHelper + })(), } } - async function authenticate( - model: Awaited<ReturnType<typeof validateModel>>, - providerInfo: Awaited<ReturnType<typeof selectProvider>>, - ) { + async function authenticate(modelInfo: ModelInfo, providerInfo: ProviderInfo) { const apiKey = opts.parseApiKey(input.request.headers) if (!apiKey) { - if (model.allowAnonymous) return + if (modelInfo.allowAnonymous) return throw new AuthError("Missing API key.") } @@ -282,7 +297,7 @@ export async function handler( .innerJoin(WorkspaceTable, eq(WorkspaceTable.id, KeyTable.workspaceID)) .innerJoin(BillingTable, eq(BillingTable.workspaceID, KeyTable.workspaceID)) .innerJoin(UserTable, and(eq(UserTable.workspaceID, KeyTable.workspaceID), eq(UserTable.id, KeyTable.userID))) - .leftJoin(ModelTable, and(eq(ModelTable.workspaceID, KeyTable.workspaceID), eq(ModelTable.model, model.id))) + .leftJoin(ModelTable, and(eq(ModelTable.workspaceID, KeyTable.workspaceID), eq(ModelTable.model, modelInfo.id))) .leftJoin( ProviderTable, and(eq(ProviderTable.workspaceID, KeyTable.workspaceID), eq(ProviderTable.provider, providerInfo.id)), @@ -308,11 +323,11 @@ export async function handler( } } - function validateBilling(model: Model, authInfo: Awaited<ReturnType<typeof authenticate>>) { + function validateBilling(authInfo: AuthInfo, modelInfo: ModelInfo) { if (!authInfo) return if (authInfo.provider?.credentials) return if (authInfo.isFree) return - if (model.allowAnonymous) return + if (modelInfo.allowAnonymous) return const billing = authInfo.billing if (!billing.paymentMethodID) @@ -356,26 +371,18 @@ export async function handler( } } - function validateModelSettings(authInfo: Awaited<ReturnType<typeof authenticate>>) { + function validateModelSettings(authInfo: AuthInfo) { if (!authInfo) return if (authInfo.isDisabled) throw new ModelError("Model is disabled") } - function updateProviderKey( - authInfo: Awaited<ReturnType<typeof authenticate>>, - providerInfo: Awaited<ReturnType<typeof selectProvider>>, - ) { + function updateProviderKey(authInfo: AuthInfo, providerInfo: ProviderInfo) { if (!authInfo) return if (!authInfo.provider?.credentials) return providerInfo.apiKey = authInfo.provider.credentials } - async function trackUsage( - authInfo: Awaited<ReturnType<typeof authenticate>>, - modelInfo: ReturnType<typeof validateModel>, - providerInfo: Awaited<ReturnType<typeof selectProvider>>, - usage: any, - ) { + async function trackUsage(authInfo: AuthInfo, modelInfo: ModelInfo, providerInfo: ProviderInfo, usage: any) { const { inputTokens, outputTokens, reasoningTokens, cacheReadTokens, cacheWrite5mTokens, cacheWrite1hTokens } = providerInfo.normalizeUsage(usage) @@ -483,7 +490,7 @@ export async function handler( ) } - async function reload(authInfo: Awaited<ReturnType<typeof authenticate>>) { + async function reload(authInfo: AuthInfo) { if (!authInfo) return if (authInfo.isFree) return if (authInfo.provider?.credentials) return diff --git a/packages/console/app/src/routes/zen/util/rateLimiter.ts b/packages/console/app/src/routes/zen/util/rateLimiter.ts new file mode 100644 index 000000000..84cd5253e --- /dev/null +++ b/packages/console/app/src/routes/zen/util/rateLimiter.ts @@ -0,0 +1,36 @@ +import { Resource } from "@opencode-ai/console-resource" +import { RateLimitError } from "./error" +import { logger } from "./logger" + +export function createRateLimiter(model: string, limit: number | undefined, ip: string) { + if (!limit) return + + const now = Date.now() + const currKey = `usage:${ip}:${model}:${buildYYYYMMDDHH(now)}` + const prevKey = `usage:${ip}:${model}:${buildYYYYMMDDHH(now - 3_600_000)}` + let currRate: number + let prevRate: number + + return { + track: async () => { + await Resource.GatewayKv.put(currKey, currRate + 1, { expirationTtl: 3600 }) + }, + check: async () => { + const values = await Resource.GatewayKv.get([currKey, prevKey]) + const prevValue = values?.get(prevKey) + const currValue = values?.get(currKey) + prevRate = prevValue ? parseInt(prevValue) : 0 + currRate = currValue ? parseInt(currValue) : 0 + logger.debug(`rate limit ${model} prev/curr: ${prevRate}/${currRate}`) + if (prevRate + currRate >= limit) + throw new RateLimitError(`Rate limit exceeded. Please try again later.`) + }, + } +} + +function buildYYYYMMDDHH(timestamp: number) { + return new Date(timestamp) + .toISOString() + .replace(/[^0-9]/g, "") + .substring(0, 10) +} |