mirror of
https://github.com/anomalyco/opencode.git
synced 2026-05-05 16:02:19 +00:00
zen: rate limiter
This commit is contained in:
parent
2c36bf9490
commit
8df7ccc304
12 changed files with 7757 additions and 23 deletions
|
|
@ -21,6 +21,7 @@ import {
|
|||
MonthlyLimitError,
|
||||
UserLimitError,
|
||||
ModelError,
|
||||
RateLimitError,
|
||||
FreeUsageLimitError,
|
||||
SubscriptionUsageLimitError,
|
||||
} from "./error"
|
||||
|
|
@ -35,7 +36,8 @@ import { anthropicHelper } from "./provider/anthropic"
|
|||
import { googleHelper } from "./provider/google"
|
||||
import { openaiHelper } from "./provider/openai"
|
||||
import { oaCompatHelper } from "./provider/openai-compatible"
|
||||
import { createRateLimiter } from "./rateLimiter"
|
||||
import { createRateLimiter as createIpRateLimiter } from "./ipRateLimiter"
|
||||
import { createRateLimiter as createKeyRateLimiter } from "./keyRateLimiter"
|
||||
import { createDataDumper } from "./dataDumper"
|
||||
import { createTrialLimiter } from "./trialLimiter"
|
||||
import { createStickyTracker } from "./stickyProviderTracker"
|
||||
|
|
@ -92,6 +94,8 @@ export async function handler(
|
|||
const isStream = opts.parseIsStream(url, body)
|
||||
const rawIp = input.request.headers.get("x-real-ip") ?? ""
|
||||
const ip = rawIp.includes(":") ? rawIp.split(":").slice(0, 4).join(":") : rawIp
|
||||
const rawZenApiKey = opts.parseApiKey(input.request.headers)
|
||||
const zenApiKey = rawZenApiKey === "public" ? undefined : rawZenApiKey
|
||||
const sessionId = input.request.headers.get("x-opencode-session") ?? ""
|
||||
const requestId = input.request.headers.get("x-opencode-request") ?? ""
|
||||
const projectId = input.request.headers.get("x-opencode-project") ?? ""
|
||||
|
|
@ -108,17 +112,13 @@ export async function handler(
|
|||
const dataDumper = createDataDumper(sessionId, requestId, projectId)
|
||||
const trialLimiter = createTrialLimiter(modelInfo.trialProvider, ip)
|
||||
const trialProviders = await trialLimiter?.check()
|
||||
const rateLimiter = createRateLimiter(
|
||||
modelInfo.id,
|
||||
modelInfo.allowAnonymous,
|
||||
modelInfo.rateLimit,
|
||||
ip,
|
||||
input.request,
|
||||
)
|
||||
const rateLimiter = modelInfo.allowAnonymous
|
||||
? createIpRateLimiter(modelInfo.id, modelInfo.rateLimit, ip, input.request)
|
||||
: createKeyRateLimiter(modelInfo.id, zenApiKey, input.request)
|
||||
await rateLimiter?.check()
|
||||
const stickyTracker = createStickyTracker(modelInfo.stickyProvider, sessionId)
|
||||
const stickyProvider = await stickyTracker?.get()
|
||||
const authInfo = await authenticate(modelInfo)
|
||||
const authInfo = await authenticate(modelInfo, zenApiKey)
|
||||
const billingSource = validateBilling(authInfo, modelInfo)
|
||||
logger.metric({ source: billingSource })
|
||||
|
||||
|
|
@ -363,7 +363,11 @@ export async function handler(
|
|||
{ status: 401 },
|
||||
)
|
||||
|
||||
if (error instanceof FreeUsageLimitError || error instanceof SubscriptionUsageLimitError) {
|
||||
if (
|
||||
error instanceof RateLimitError ||
|
||||
error instanceof FreeUsageLimitError ||
|
||||
error instanceof SubscriptionUsageLimitError
|
||||
) {
|
||||
const headers = new Headers()
|
||||
if (error.retryAfter) {
|
||||
headers.set("retry-after", String(error.retryAfter))
|
||||
|
|
@ -492,9 +496,8 @@ export async function handler(
|
|||
}
|
||||
}
|
||||
|
||||
async function authenticate(modelInfo: ModelInfo) {
|
||||
const apiKey = opts.parseApiKey(input.request.headers)
|
||||
if (!apiKey || apiKey === "public") {
|
||||
async function authenticate(modelInfo: ModelInfo, zenApiKey?: string) {
|
||||
if (!zenApiKey) {
|
||||
if (modelInfo.allowAnonymous) return
|
||||
throw new AuthError(t("zen.api.error.missingApiKey"))
|
||||
}
|
||||
|
|
@ -573,7 +576,7 @@ export async function handler(
|
|||
isNull(LiteTable.timeDeleted),
|
||||
),
|
||||
)
|
||||
.where(and(eq(KeyTable.key, apiKey), isNull(KeyTable.timeDeleted)))
|
||||
.where(and(eq(KeyTable.key, zenApiKey), isNull(KeyTable.timeDeleted)))
|
||||
.then((rows) => rows[0]),
|
||||
)
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue