zen: tpm routing

1 files changed, 21 insertions, 12 deletions

diff --git a/packages/console/app/src/routes/zen/util/modelTpmLimiter.ts b/packages/console/app/src/routes/zen/util/modelTpmLimiter.ts
index 1fdf360ed..53015d51c 100644
--- a/packages/console/app/src/routes/zen/util/modelTpmLimiter.ts
+++ b/packages/console/app/src/routes/zen/util/modelTpmLimiter.ts
@@ -1,10 +1,10 @@
 import { and, Database, eq, inArray, sql } from "@opencode-ai/console-core/drizzle/index.js"
-import { ModelRateLimitTable } from "@opencode-ai/console-core/schema/ip.sql.js"
+import { ModelTpmLimitTable } from "@opencode-ai/console-core/schema/ip.sql.js"
 import { UsageInfo } from "./provider/provider"
 
 export function createModelTpmLimiter(providers: { id: string; model: string; tpmLimit?: number }[]) {
-  const keys = providers.filter((p) => p.tpmLimit).map((p) => `${p.id}/${p.model}`)
-  if (keys.length === 0) return
+  const ids = providers.filter((p) => p.tpmLimit).map((p) => `${p.id}/${p.model}`)
+  if (ids.length === 0) return
 
   const yyyyMMddHHmm = new Date(Date.now())
     .toISOString()
@@ -16,30 +16,39 @@ export function createModelTpmLimiter(providers: { id: string; model: string; tp
       const data = await Database.use((tx) =>
         tx
           .select()
-          .from(ModelRateLimitTable)
-          .where(and(inArray(ModelRateLimitTable.key, keys), eq(ModelRateLimitTable.interval, yyyyMMddHHmm))),
+          .from(ModelTpmLimitTable)
+          .where(
+            inArray(
+              ModelTpmLimitTable.id,
+              ids.map((id) => formatId(id, yyyyMMddHHmm)),
+            ),
+          ),
       )
 
       // convert to map of model to count
       return data.reduce(
         (acc, curr) => {
-          acc[curr.key] = curr.count
+          acc[curr.id] = curr.count
           return acc
         },
         {} as Record<string, number>,
       )
     },
-    track: async (id: string, model: string, usageInfo: UsageInfo) => {
-      const key = `${id}/${model}`
-      if (!keys.includes(key)) return
+    track: async (provider: string, model: string, usageInfo: UsageInfo) => {
+      const id = `${provider}/${model}`
+      if (!ids.includes(id)) return
       const usage = usageInfo.inputTokens
       if (usage <= 0) return
       await Database.use((tx) =>
         tx
-          .insert(ModelRateLimitTable)
-          .values({ key, interval: yyyyMMddHHmm, count: usage })
-          .onDuplicateKeyUpdate({ set: { count: sql`${ModelRateLimitTable.count} + ${usage}` } }),
+          .insert(ModelTpmLimitTable)
+          .values({ id: formatId(id, yyyyMMddHHmm), count: usage })
+          .onDuplicateKeyUpdate({ set: { count: sql`${ModelTpmLimitTable.count} + ${usage}` } }),
       )
     },
   }
+
+  function formatId(id: string, yyyyMMddHHmm: string) {
+    return `${id.substring(0, 200)}/${yyyyMMddHHmm}`
+  }
 }