zen: trial

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

@@ -13,13 +13,20 @@ import { ModelTable } from "@opencode-ai/console-core/schema/model.sql.js"
 import { ProviderTable } from "@opencode-ai/console-core/schema/provider.sql.js"
 import { logger } from "./logger"
 import { AuthError, CreditsError, MonthlyLimitError, UserLimitError, ModelError, RateLimitError } from "./error"
-import { createBodyConverter, createStreamPartConverter, createResponseConverter } from "./provider/provider"
+import {
+  createBodyConverter,
+  createStreamPartConverter,
+  createResponseConverter,
+  ProviderHelper,
+  UsageInfo,
+} from "./provider/provider"
 import { anthropicHelper } from "./provider/anthropic"
 import { googleHelper } from "./provider/google"
 import { openaiHelper } from "./provider/openai"
 import { oaCompatHelper } from "./provider/openai-compatible"
 import { createRateLimiter } from "./rateLimiter"
 import { createDataDumper } from "./dataDumper"
+import { createTrialLimiter } from "./trialLimiter"
 
 type ZenData = Awaited<ReturnType<typeof ZenData.list>>
 type RetryOptions = {
@@ -62,11 +69,13 @@ export async function handler(
     const zenData = ZenData.list()
     const modelInfo = validateModel(zenData, model)
     const dataDumper = createDataDumper(sessionId, requestId)
+    const trialLimiter = createTrialLimiter(modelInfo.trial?.limit, ip)
+    const isTrial = await trialLimiter?.isTrial()
     const rateLimiter = createRateLimiter(modelInfo.id, modelInfo.rateLimit, ip)
     await rateLimiter?.check()
 
     const retriableRequest = async (retry: RetryOptions = { excludeProviders: [], retryCount: 0 }) => {
-      const providerInfo = selectProvider(zenData, modelInfo, sessionId, retry)
+      const providerInfo = selectProvider(zenData, modelInfo, sessionId, isTrial ?? false, retry)
       const authInfo = await authenticate(modelInfo, providerInfo)
       validateBilling(authInfo, modelInfo)
       validateModelSettings(authInfo)
@@ -136,8 +145,10 @@ export async function handler(
       logger.debug("RESPONSE: " + body)
       dataDumper?.provideResponse(body)
       dataDumper?.flush()
+      const tokensInfo = providerInfo.normalizeUsage(json.usage)
+      await trialLimiter?.track(tokensInfo)
       await rateLimiter?.track()
-      await trackUsage(authInfo, modelInfo, providerInfo, json.usage)
+      await trackUsage(authInfo, modelInfo, providerInfo, tokensInfo)
       await reload(authInfo)
       return new Response(body, {
         status: res.status,
@@ -169,7 +180,9 @@ export async function handler(
                 await rateLimiter?.track()
                 const usage = usageParser.retrieve()
                 if (usage) {
-                  await trackUsage(authInfo, modelInfo, providerInfo, usage)
+                  const tokensInfo = providerInfo.normalizeUsage(usage)
+                  await trialLimiter?.track(tokensInfo)
+                  await trackUsage(authInfo, modelInfo, providerInfo, tokensInfo)
                   await reload(authInfo)
                 }
                 c.close()
@@ -275,8 +288,19 @@ export async function handler(
     return { id: modelId, ...modelData }
   }
 
-  function selectProvider(zenData: ZenData, modelInfo: ModelInfo, sessionId: string, retry: RetryOptions) {
+  function selectProvider(
+    zenData: ZenData,
+    modelInfo: ModelInfo,
+    sessionId: string,
+    isTrial: boolean,
+    retry: RetryOptions,
+  ) {
     const provider = (() => {
+      // temporarily commment out
+      //if (isTrial) {
+      //  return modelInfo.providers.find((provider) => provider.id === modelInfo.trial!.provider)
+      //}
+
       if (retry.retryCount === MAX_RETRIES) {
         return modelInfo.providers.find((provider) => provider.id === modelInfo.fallbackProvider)
       }
@@ -432,9 +456,14 @@ export async function handler(
     providerInfo.apiKey = authInfo.provider.credentials
   }
 
-  async function trackUsage(authInfo: AuthInfo, modelInfo: ModelInfo, providerInfo: ProviderInfo, usage: any) {
+  async function trackUsage(
+    authInfo: AuthInfo,
+    modelInfo: ModelInfo,
+    providerInfo: ProviderInfo,
+    usageInfo: UsageInfo,
+  ) {
     const { inputTokens, outputTokens, reasoningTokens, cacheReadTokens, cacheWrite5mTokens, cacheWrite1hTokens } =
-      providerInfo.normalizeUsage(usage)
+      usageInfo
 
     const modelCost =
       modelInfo.cost200K &&

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}` } }),
+      )
+    },
+  }
+}

packages/console/core/migrations/0038_famous_magik.sql

@@ -0,0 +1,8 @@
+CREATE TABLE `ip` (
+	`ip` varchar(45) NOT NULL,
+	`time_created` timestamp(3) NOT NULL DEFAULT (now()),
+	`time_updated` timestamp(3) NOT NULL DEFAULT CURRENT_TIMESTAMP(3) ON UPDATE CURRENT_TIMESTAMP(3),
+	`time_deleted` timestamp(3),
+	`usage` int,
+	CONSTRAINT `ip_ip_pk` PRIMARY KEY(`ip`)
+);

packages/console/core/migrations/meta/_journal.json

@@ -267,6 +267,13 @@
       "when": 1761928273807,
       "tag": "0037_messy_jackal",
       "breakpoints": true
+    },
+    {
+      "idx": 38,
+      "version": "5",
+      "when": 1764110043942,
+      "tag": "0038_famous_magik",
+      "breakpoints": true
     }
   ]
-}
+}

packages/console/core/src/model.ts

@@ -24,6 +24,12 @@ export namespace ZenData {
     cost: ModelCostSchema,
     cost200K: ModelCostSchema.optional(),
     allowAnonymous: z.boolean().optional(),
+    trial: z
+      .object({
+        limit: z.number(),
+        provider: z.string(),
+      })
+      .optional(),
     rateLimit: z.number().optional(),
     fallbackProvider: z.string().optional(),
     providers: z.array(

packages/console/core/src/schema/ip.sql.ts

@@ -0,0 +1,12 @@
+import { mysqlTable, int, primaryKey, varchar } from "drizzle-orm/mysql-core"
+import { timestamps } from "../drizzle/types"
+
+export const IpTable = mysqlTable(
+  "ip",
+  {
+    ip: varchar("ip", { length: 45 }).notNull(),
+    ...timestamps,
+    usage: int("usage"),
+  },
+  (table) => [primaryKey({ columns: [table.ip] })],
+)