mirror of
https://gitea.toothfairyai.com/ToothFairyAI/tf_code.git
synced 2026-04-12 19:59:00 +00:00
wip: zen
This commit is contained in:
@@ -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) {
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user