43 lines
1.2 KiB
TypeScript
43 lines
1.2 KiB
TypeScript
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"
|
|
|
|
export function createTrialLimiter(limit: number | undefined, ip: string) {
|
|
if (!limit) return
|
|
|
|
let trial: boolean
|
|
|
|
return {
|
|
isTrial: async () => {
|
|
const data = await Database.use((tx) =>
|
|
tx
|
|
.select({
|
|
usage: IpTable.usage,
|
|
})
|
|
.from(IpTable)
|
|
.where(eq(IpTable.ip, ip))
|
|
.then((rows) => rows[0]),
|
|
)
|
|
|
|
trial = (data?.usage ?? 0) < limit
|
|
return trial
|
|
},
|
|
track: async (usageInfo: UsageInfo) => {
|
|
if (!trial) 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}` } }),
|
|
)
|
|
},
|
|
}
|
|
}
|