This commit is contained in:
Frank
2026-03-03 00:25:03 -05:00
parent 9f150b0776
commit 6aa4928e9e
23 changed files with 422 additions and 313 deletions

View File

@@ -97,9 +97,9 @@ export async function handler(
const zenData = ZenData.list(opts.modelList)
const modelInfo = validateModel(zenData, model)
const dataDumper = createDataDumper(sessionId, requestId, projectId)
const trialLimiter = createTrialLimiter(modelInfo.trial, ip, ocClient)
const isTrial = await trialLimiter?.isTrial()
const rateLimiter = createRateLimiter(modelInfo.rateLimit, ip, input.request)
const trialLimiter = createTrialLimiter(modelInfo.trialProvider, ip)
const trialProvider = await trialLimiter?.check()
const rateLimiter = createRateLimiter(modelInfo.allowAnonymous, ip, input.request)
await rateLimiter?.check()
const stickyTracker = createStickyTracker(modelInfo.stickyProvider, sessionId)
const stickyProvider = await stickyTracker?.get()
@@ -114,7 +114,7 @@ export async function handler(
authInfo,
modelInfo,
sessionId,
isTrial ?? false,
trialProvider,
retry,
stickyProvider,
)
@@ -144,9 +144,6 @@ export async function handler(
Object.entries(providerInfo.headerMappings ?? {}).forEach(([k, v]) => {
headers.set(k, headers.get(v)!)
})
Object.entries(providerInfo.headers ?? {}).forEach(([k, v]) => {
headers.set(k, v)
})
headers.delete("host")
headers.delete("content-length")
headers.delete("x-opencode-request")
@@ -295,18 +292,13 @@ export async function handler(
part = part.trim()
usageParser.parse(part)
if (providerInfo.responseModifier) {
for (const [k, v] of Object.entries(providerInfo.responseModifier)) {
part = part.replace(k, v)
}
c.enqueue(encoder.encode(part + "\n\n"))
} else if (providerInfo.format !== opts.format) {
if (providerInfo.format !== opts.format) {
part = streamConverter(part)
c.enqueue(encoder.encode(part + "\n\n"))
}
}
if (!providerInfo.responseModifier && providerInfo.format === opts.format) {
if (providerInfo.format === opts.format) {
c.enqueue(value)
}
@@ -398,7 +390,7 @@ export async function handler(
authInfo: AuthInfo,
modelInfo: ModelInfo,
sessionId: string,
isTrial: boolean,
trialProvider: string | undefined,
retry: RetryOptions,
stickyProvider: string | undefined,
) {
@@ -407,8 +399,8 @@ export async function handler(
return modelInfo.providers.find((provider) => provider.id === modelInfo.byokProvider)
}
if (isTrial) {
return modelInfo.providers.find((provider) => provider.id === modelInfo.trial!.provider)
if (trialProvider) {
return modelInfo.providers.find((provider) => provider.id === trialProvider)
}
if (stickyProvider) {

View File

@@ -2,29 +2,28 @@ import { Database, eq, and, sql, inArray } from "@opencode-ai/console-core/drizz
import { IpRateLimitTable } from "@opencode-ai/console-core/schema/ip.sql.js"
import { FreeUsageLimitError } from "./error"
import { logger } from "./logger"
import { ZenData } from "@opencode-ai/console-core/model.js"
import { i18n } from "~/i18n"
import { localeFromRequest } from "~/lib/language"
import { Subscription } from "@opencode-ai/console-core/subscription.js"
export function createRateLimiter(limit: ZenData.RateLimit | undefined, rawIp: string, request: Request) {
if (!limit) return
export function createRateLimiter(allowAnonymous: boolean | undefined, rawIp: string, request: Request) {
if (!allowAnonymous) return
const dict = i18n(localeFromRequest(request))
const limitValue = limit.checkHeader && !request.headers.get(limit.checkHeader) ? limit.fallbackValue! : limit.value
const limits = Subscription.getFreeLimits()
const limitValue =
limits.checkHeader && !request.headers.get(limits.checkHeader) ? limits.fallbackValue : limits.dailyRequests
const ip = !rawIp.length ? "unknown" : rawIp
const now = Date.now()
const intervals =
limit.period === "day"
? [buildYYYYMMDD(now)]
: [buildYYYYMMDDHH(now), buildYYYYMMDDHH(now - 3_600_000), buildYYYYMMDDHH(now - 7_200_000)]
const interval = buildYYYYMMDD(now)
return {
track: async () => {
await Database.use((tx) =>
tx
.insert(IpRateLimitTable)
.values({ ip, interval: intervals[0], count: 1 })
.values({ ip, interval, count: 1 })
.onDuplicateKeyUpdate({ set: { count: sql`${IpRateLimitTable.count} + 1` } }),
)
},
@@ -33,15 +32,12 @@ export function createRateLimiter(limit: ZenData.RateLimit | undefined, rawIp: s
tx
.select({ interval: IpRateLimitTable.interval, count: IpRateLimitTable.count })
.from(IpRateLimitTable)
.where(and(eq(IpRateLimitTable.ip, ip), inArray(IpRateLimitTable.interval, intervals))),
.where(and(eq(IpRateLimitTable.ip, ip), inArray(IpRateLimitTable.interval, [interval]))),
)
const total = rows.reduce((sum, r) => sum + r.count, 0)
logger.debug(`rate limit total: ${total}`)
if (total >= limitValue)
throw new FreeUsageLimitError(
dict["zen.api.error.rateLimitExceeded"],
limit.period === "day" ? getRetryAfterDay(now) : getRetryAfterHour(rows, intervals, limitValue, now),
)
throw new FreeUsageLimitError(dict["zen.api.error.rateLimitExceeded"], getRetryAfterDay(now))
},
}
}
@@ -50,37 +46,9 @@ 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) {
return new Date(timestamp)
.toISOString()
.replace(/[^0-9]/g, "")
.substring(0, 8)
}
function buildYYYYMMDDHH(timestamp: number) {
return new Date(timestamp)
.toISOString()
.replace(/[^0-9]/g, "")
.substring(0, 10)
}

View File

@@ -1,21 +1,18 @@
import { Database, eq, sql } from "@opencode-ai/console-core/drizzle/index.js"
import { IpTable } from "@opencode-ai/console-core/schema/ip.sql.js"
import { UsageInfo } from "./provider/provider"
import { ZenData } from "@opencode-ai/console-core/model.js"
import { Subscription } from "@opencode-ai/console-core/subscription.js"
export function createTrialLimiter(trial: ZenData.Trial | undefined, ip: string, client: string) {
if (!trial) return
export function createTrialLimiter(trialProvider: string | undefined, ip: string) {
if (!trialProvider) return
if (!ip) return
const limit =
trial.limits.find((limit) => limit.client === client)?.limit ??
trial.limits.find((limit) => limit.client === undefined)?.limit
if (!limit) return
const limit = Subscription.getFreeLimits().promoTokens
let _isTrial: boolean
return {
isTrial: async () => {
check: async () => {
const data = await Database.use((tx) =>
tx
.select({
@@ -27,7 +24,7 @@ export function createTrialLimiter(trial: ZenData.Trial | undefined, ip: string,
)
_isTrial = (data?.usage ?? 0) < limit
return _isTrial
return _isTrial ? trialProvider : undefined
},
track: async (usageInfo: UsageInfo) => {
if (!_isTrial) return