Skip to content

Commit 6aa4928

Browse files
fwangfwang
committed
wip: zen
1 parent 9f150b0 commit 6aa4928

23 files changed

Lines changed: 422 additions & 313 deletions

‎infra/console.ts‎

Lines changed: 1 addition & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -118,7 +118,6 @@ const ZEN_LITE_PRICE = new sst.Linkable("ZEN_LITE_PRICE", {
118118
price: zenLitePrice.id,
119119
},
120120
})
121-
const ZEN_LITE_LIMITS = new sst.Secret("ZEN_LITE_LIMITS")
122121

123122
const zenBlackProduct = new stripe.Product("ZenBlack", {
124123
name: "OpenCode Black",
@@ -142,7 +141,6 @@ const ZEN_BLACK_PRICE = new sst.Linkable("ZEN_BLACK_PRICE", {
142141
plan20: zenBlackPrice20.id,
143142
},
144143
})
145-
const ZEN_BLACK_LIMITS = new sst.Secret("ZEN_BLACK_LIMITS")
146144

147145
const ZEN_MODELS = [
148146
new sst.Secret("ZEN_MODELS1"),
@@ -215,9 +213,8 @@ new sst.cloudflare.x.SolidStart("Console", {
215213
AWS_SES_ACCESS_KEY_ID,
216214
AWS_SES_SECRET_ACCESS_KEY,
217215
ZEN_BLACK_PRICE,
218-
ZEN_BLACK_LIMITS,
219216
ZEN_LITE_PRICE,
220-
ZEN_LITE_LIMITS,
217+
new sst.Secret("ZEN_LIMITS"),
221218
new sst.Secret("ZEN_SESSION_SECRET"),
222219
...ZEN_MODELS,
223220
...($dev

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

Lines changed: 9 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -97,9 +97,9 @@ export async function handler(
9797
const zenData = ZenData.list(opts.modelList)
9898
const modelInfo = validateModel(zenData, model)
9999
const dataDumper = createDataDumper(sessionId, requestId, projectId)
100-
const trialLimiter = createTrialLimiter(modelInfo.trial, ip, ocClient)
101-
const isTrial = await trialLimiter?.isTrial()
102-
const rateLimiter = createRateLimiter(modelInfo.rateLimit, ip, input.request)
100+
const trialLimiter = createTrialLimiter(modelInfo.trialProvider, ip)
101+
const trialProvider = await trialLimiter?.check()
102+
const rateLimiter = createRateLimiter(modelInfo.allowAnonymous, ip, input.request)
103103
await rateLimiter?.check()
104104
const stickyTracker = createStickyTracker(modelInfo.stickyProvider, sessionId)
105105
const stickyProvider = await stickyTracker?.get()
@@ -114,7 +114,7 @@ export async function handler(
114114
authInfo,
115115
modelInfo,
116116
sessionId,
117-
isTrial ?? false,
117+
trialProvider,
118118
retry,
119119
stickyProvider,
120120
)
@@ -144,9 +144,6 @@ export async function handler(
144144
Object.entries(providerInfo.headerMappings ?? {}).forEach(([k, v]) => {
145145
headers.set(k, headers.get(v)!)
146146
})
147-
Object.entries(providerInfo.headers ?? {}).forEach(([k, v]) => {
148-
headers.set(k, v)
149-
})
150147
headers.delete("host")
151148
headers.delete("content-length")
152149
headers.delete("x-opencode-request")
@@ -295,18 +292,13 @@ export async function handler(
295292
part = part.trim()
296293
usageParser.parse(part)
297294

298-
if (providerInfo.responseModifier) {
299-
for (const [k, v] of Object.entries(providerInfo.responseModifier)) {
300-
part = part.replace(k, v)
301-
}
302-
c.enqueue(encoder.encode(part + "\n\n"))
303-
} else if (providerInfo.format !== opts.format) {
295+
if (providerInfo.format !== opts.format) {
304296
part = streamConverter(part)
305297
c.enqueue(encoder.encode(part + "\n\n"))
306298
}
307299
}
308300

309-
if (!providerInfo.responseModifier && providerInfo.format === opts.format) {
301+
if (providerInfo.format === opts.format) {
310302
c.enqueue(value)
311303
}
312304

@@ -398,7 +390,7 @@ export async function handler(
398390
authInfo: AuthInfo,
399391
modelInfo: ModelInfo,
400392
sessionId: string,
401-
isTrial: boolean,
393+
trialProvider: string | undefined,
402394
retry: RetryOptions,
403395
stickyProvider: string | undefined,
404396
) {
@@ -407,8 +399,8 @@ export async function handler(
407399
return modelInfo.providers.find((provider) => provider.id === modelInfo.byokProvider)
408400
}
409401

410-
if (isTrial) {
411-
return modelInfo.providers.find((provider) => provider.id === modelInfo.trial!.provider)
402+
if (trialProvider) {
403+
return modelInfo.providers.find((provider) => provider.id === trialProvider)
412404
}
413405

414406
if (stickyProvider) {

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

Lines changed: 10 additions & 42 deletions
Original file line numberDiff line numberDiff line change
@@ -2,29 +2,28 @@ import { Database, eq, and, sql, inArray } from "@opencode-ai/console-core/drizz
22
import { IpRateLimitTable } from "@opencode-ai/console-core/schema/ip.sql.js"
33
import { FreeUsageLimitError } from "./error"
44
import { logger } from "./logger"
5-
import { ZenData } from "@opencode-ai/console-core/model.js"
65
import { i18n } from "~/i18n"
76
import { localeFromRequest } from "~/lib/language"
7+
import { Subscription } from "@opencode-ai/console-core/subscription.js"
88

9-
export function createRateLimiter(limit: ZenData.RateLimit | undefined, rawIp: string, request: Request) {
10-
if (!limit) return
9+
export function createRateLimiter(allowAnonymous: boolean | undefined, rawIp: string, request: Request) {
10+
if (!allowAnonymous) return
1111
const dict = i18n(localeFromRequest(request))
1212

13-
const limitValue = limit.checkHeader && !request.headers.get(limit.checkHeader) ? limit.fallbackValue! : limit.value
13+
const limits = Subscription.getFreeLimits()
14+
const limitValue =
15+
limits.checkHeader && !request.headers.get(limits.checkHeader) ? limits.fallbackValue : limits.dailyRequests
1416

1517
const ip = !rawIp.length ? "unknown" : rawIp
1618
const now = Date.now()
17-
const intervals =
18-
limit.period === "day"
19-
? [buildYYYYMMDD(now)]
20-
: [buildYYYYMMDDHH(now), buildYYYYMMDDHH(now - 3_600_000), buildYYYYMMDDHH(now - 7_200_000)]
19+
const interval = buildYYYYMMDD(now)
2120

2221
return {
2322
track: async () => {
2423
await Database.use((tx) =>
2524
tx
2625
.insert(IpRateLimitTable)
27-
.values({ ip, interval: intervals[0], count: 1 })
26+
.values({ ip, interval, count: 1 })
2827
.onDuplicateKeyUpdate({ set: { count: sql`${IpRateLimitTable.count} + 1` } }),
2928
)
3029
},
@@ -33,15 +32,12 @@ export function createRateLimiter(limit: ZenData.RateLimit | undefined, rawIp: s
3332
tx
3433
.select({ interval: IpRateLimitTable.interval, count: IpRateLimitTable.count })
3534
.from(IpRateLimitTable)
36-
.where(and(eq(IpRateLimitTable.ip, ip), inArray(IpRateLimitTable.interval, intervals))),
35+
.where(and(eq(IpRateLimitTable.ip, ip), inArray(IpRateLimitTable.interval, [interval]))),
3736
)
3837
const total = rows.reduce((sum, r) => sum + r.count, 0)
3938
logger.debug(`rate limit total: ${total}`)
4039
if (total >= limitValue)
41-
throw new FreeUsageLimitError(
42-
dict["zen.api.error.rateLimitExceeded"],
43-
limit.period === "day" ? getRetryAfterDay(now) : getRetryAfterHour(rows, intervals, limitValue, now),
44-
)
40+
throw new FreeUsageLimitError(dict["zen.api.error.rateLimitExceeded"], getRetryAfterDay(now))
4541
},
4642
}
4743
}
@@ -50,37 +46,9 @@ export function getRetryAfterDay(now: number) {
5046
return Math.ceil((86_400_000 - (now % 86_400_000)) / 1000)
5147
}
5248

53-
export function getRetryAfterHour(
54-
rows: { interval: string; count: number }[],
55-
intervals: string[],
56-
limit: number,
57-
now: number,
58-
) {
59-
const counts = new Map(rows.map((r) => [r.interval, r.count]))
60-
// intervals are ordered newest to oldest: [current, -1h, -2h]
61-
// simulate dropping oldest intervals one at a time
62-
let running = intervals.reduce((sum, i) => sum + (counts.get(i) ?? 0), 0)
63-
for (let i = intervals.length - 1; i >= 0; i--) {
64-
running -= counts.get(intervals[i]) ?? 0
65-
if (running < limit) {
66-
// interval at index i rolls out of the window (intervals.length - i) hours from the current hour start
67-
const hours = intervals.length - i
68-
return Math.ceil((hours * 3_600_000 - (now % 3_600_000)) / 1000)
69-
}
70-
}
71-
return Math.ceil((3_600_000 - (now % 3_600_000)) / 1000)
72-
}
73-
7449
function buildYYYYMMDD(timestamp: number) {
7550
return new Date(timestamp)
7651
.toISOString()
7752
.replace(/[^0-9]/g, "")
7853
.substring(0, 8)
7954
}
80-
81-
function buildYYYYMMDDHH(timestamp: number) {
82-
return new Date(timestamp)
83-
.toISOString()
84-
.replace(/[^0-9]/g, "")
85-
.substring(0, 10)
86-
}

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

Lines changed: 6 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -1,21 +1,18 @@
11
import { Database, eq, sql } from "@opencode-ai/console-core/drizzle/index.js"
22
import { IpTable } from "@opencode-ai/console-core/schema/ip.sql.js"
33
import { UsageInfo } from "./provider/provider"
4-
import { ZenData } from "@opencode-ai/console-core/model.js"
4+
import { Subscription } from "@opencode-ai/console-core/subscription.js"
55

6-
export function createTrialLimiter(trial: ZenData.Trial | undefined, ip: string, client: string) {
7-
if (!trial) return
6+
export function createTrialLimiter(trialProvider: string | undefined, ip: string) {
7+
if (!trialProvider) return
88
if (!ip) return
99

10-
const limit =
11-
trial.limits.find((limit) => limit.client === client)?.limit ??
12-
trial.limits.find((limit) => limit.client === undefined)?.limit
13-
if (!limit) return
10+
const limit = Subscription.getFreeLimits().promoTokens
1411

1512
let _isTrial: boolean
1613

1714
return {
18-
isTrial: async () => {
15+
check: async () => {
1916
const data = await Database.use((tx) =>
2017
tx
2118
.select({
@@ -27,7 +24,7 @@ export function createTrialLimiter(trial: ZenData.Trial | undefined, ip: string,
2724
)
2825

2926
_isTrial = (data?.usage ?? 0) < limit
30-
return _isTrial
27+
return _isTrial ? trialProvider : undefined
3128
},
3229
track: async (usageInfo: UsageInfo) => {
3330
if (!_isTrial) return

‎packages/console/app/test/rateLimiter.test.ts‎

Lines changed: 1 addition & 74 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
import { describe, expect, test } from "bun:test"
2-
import { getRetryAfterDay, getRetryAfterHour } from "../src/routes/zen/util/rateLimiter"
2+
import { getRetryAfterDay } from "../src/routes/zen/util/rateLimiter"
33

44
describe("getRetryAfterDay", () => {
55
test("returns full day at midnight UTC", () => {
@@ -17,76 +17,3 @@ describe("getRetryAfterDay", () => {
1717
expect(getRetryAfterDay(almost)).toBe(1)
1818
})
1919
})
20-
21-
describe("getRetryAfterHour", () => {
22-
// 14:30:00 UTC — 30 minutes into the current hour
23-
const now = Date.UTC(2026, 0, 15, 14, 30, 0, 0)
24-
const intervals = ["2026011514", "2026011513", "2026011512"]
25-
26-
test("waits 3 hours when all usage is in current hour", () => {
27-
const rows = [{ interval: "2026011514", count: 10 }]
28-
// only current hour has usage — it won't leave the window for 3 hours from hour start
29-
// 3 * 3600 - 1800 = 9000s
30-
expect(getRetryAfterHour(rows, intervals, 10, now)).toBe(9000)
31-
})
32-
33-
test("waits 1 hour when dropping oldest interval is sufficient", () => {
34-
const rows = [
35-
{ interval: "2026011514", count: 2 },
36-
{ interval: "2026011512", count: 10 },
37-
]
38-
// total=12, drop oldest (-2h, count=10) -> 2 < 10
39-
// hours = 3 - 2 = 1 -> 1 * 3600 - 1800 = 1800s
40-
expect(getRetryAfterHour(rows, intervals, 10, now)).toBe(1800)
41-
})
42-
43-
test("waits 2 hours when usage spans oldest two intervals", () => {
44-
const rows = [
45-
{ interval: "2026011513", count: 8 },
46-
{ interval: "2026011512", count: 5 },
47-
]
48-
// total=13, drop -2h (5) -> 8, 8 >= 8, drop -1h (8) -> 0 < 8
49-
// hours = 3 - 1 = 2 -> 2 * 3600 - 1800 = 5400s
50-
expect(getRetryAfterHour(rows, intervals, 8, now)).toBe(5400)
51-
})
52-
53-
test("waits 1 hour when oldest interval alone pushes over limit", () => {
54-
const rows = [
55-
{ interval: "2026011514", count: 1 },
56-
{ interval: "2026011513", count: 1 },
57-
{ interval: "2026011512", count: 10 },
58-
]
59-
// total=12, drop -2h (10) -> 2 < 10
60-
// hours = 3 - 2 = 1 -> 1800s
61-
expect(getRetryAfterHour(rows, intervals, 10, now)).toBe(1800)
62-
})
63-
64-
test("waits 2 hours when middle interval keeps total over limit", () => {
65-
const rows = [
66-
{ interval: "2026011514", count: 4 },
67-
{ interval: "2026011513", count: 4 },
68-
{ interval: "2026011512", count: 4 },
69-
]
70-
// total=12, drop -2h (4) -> 8, 8 >= 5, drop -1h (4) -> 4 < 5
71-
// hours = 3 - 1 = 2 -> 5400s
72-
expect(getRetryAfterHour(rows, intervals, 5, now)).toBe(5400)
73-
})
74-
75-
test("rounds up to nearest second", () => {
76-
const offset = Date.UTC(2026, 0, 15, 14, 30, 0, 500)
77-
const rows = [
78-
{ interval: "2026011514", count: 2 },
79-
{ interval: "2026011512", count: 10 },
80-
]
81-
// hours=1 -> 3_600_000 - 1_800_500 = 1_799_500ms -> ceil(1799.5) = 1800
82-
expect(getRetryAfterHour(rows, intervals, 10, offset)).toBe(1800)
83-
})
84-
85-
test("fallback returns time until next hour when rows are empty", () => {
86-
// edge case: rows empty but function called (shouldn't happen in practice)
87-
// loop drops all zeros, running stays 0 which is < any positive limit on first iteration
88-
const rows: { interval: string; count: number }[] = []
89-
// drop -2h (0) -> 0 < 1 -> hours = 3 - 2 = 1 -> 1800s
90-
expect(getRetryAfterHour(rows, intervals, 1, now)).toBe(1800)
91-
})
92-
})

‎packages/console/core/package.json‎

Lines changed: 3 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -34,12 +34,9 @@
3434
"promote-models-to-prod": "script/promote-models.ts production",
3535
"pull-models-from-dev": "script/pull-models.ts dev",
3636
"pull-models-from-prod": "script/pull-models.ts production",
37-
"update-black": "script/update-black.ts",
38-
"promote-black-to-dev": "script/promote-black.ts dev",
39-
"promote-black-to-prod": "script/promote-black.ts production",
40-
"update-lite": "script/update-lite.ts",
41-
"promote-lite-to-dev": "script/promote-lite.ts dev",
42-
"promote-lite-to-prod": "script/promote-lite.ts production",
37+
"update-limits": "script/update-limits.ts",
38+
"promote-limits-to-dev": "script/promote-limits.ts dev",
39+
"promote-limits-to-prod": "script/promote-limits.ts production",
4340
"typecheck": "tsgo --noEmit"
4441
},
4542
"devDependencies": {

0 commit comments

Comments
 (0)