zen: trial

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

@@ -0,0 +1,42 @@
+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}` } }),
+      )
+    },
+  }
+}