blob: 319825dd79f0d8cab51282e3eb2114825c4e771a (
plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
|
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 { Subscription } from "@opencode-ai/console-core/subscription.js"
export function createTrialLimiter(trialProviders: string[] | undefined, ip: string) {
if (!trialProviders) return
if (!ip) return
const limit = Subscription.getFreeLimits().promoTokens
let _isTrial: boolean
return {
check: async () => {
const data = await Database.use((tx) =>
tx
.select({
usage: IpTable.usage,
})
.from(IpTable)
.where(eq(IpTable.ip, ip))
.then((rows) => rows[0]),
)
_isTrial = (data?.usage ?? 0) < limit
return _isTrial ? trialProviders : undefined
},
track: async (usageInfo: UsageInfo) => {
if (!_isTrial) return
const usage =
usageInfo.inputTokens +
usageInfo.outputTokens +
(usageInfo.reasoningTokens ?? 0) +
(usageInfo.cacheReadTokens ?? 0) +
(usageInfo.cacheWrite5mTokens ?? 0) +
(usageInfo.cacheWrite1hTokens ?? 0)
await Database.use((tx) =>
tx
.insert(IpTable)
.values({ ip, usage })
.onDuplicateKeyUpdate({ set: { usage: sql`${IpTable.usage} + ${usage}` } }),
)
},
}
}
|
