zen: tpm routing

2 files changed, 28 insertions, 26 deletions

diff --git a/packages/console/app/src/routes/zen/util/handler.ts b/packages/console/app/src/routes/zen/util/handler.ts
index 81c512b99..635eadebe 100644
--- a/packages/console/app/src/routes/zen/util/handler.ts
+++ b/packages/console/app/src/routes/zen/util/handler.ts
@@ -448,31 +448,40 @@ export async function handler(
         return modelInfo.providers.find((provider) => provider.id === modelInfo.byokProvider)
       }
 
+      // Filter out TPM limited providers
+      const allProviders = modelInfo.providers.filter((provider) => {
+        if (!provider.tpmLimit) return true
+        const usage = modelTpmLimits?.[`${provider.id}/${provider.model}`] ?? 0
+        return usage < provider.tpmLimit * 1_000_000
+      })
+
       // Always use the same provider for the same session
       if (stickyProvider) {
-        const provider = modelInfo.providers.find((provider) => provider.id === stickyProvider)
+        const provider = allProviders.find((provider) => provider.id === stickyProvider)
         if (provider) return provider
       }
 
       if (trialProviders) {
         const trialProvider = trialProviders[Math.floor(Math.random() * trialProviders.length)]
-        const provider = modelInfo.providers.find((provider) => provider.id === trialProvider)
+        const provider = allProviders.find((provider) => provider.id === trialProvider)
         if (provider) return provider
       }
 
       if (retry.retryCount !== MAX_FAILOVER_RETRIES) {
-        const allProviders = modelInfo.providers
+        let topPriority = Infinity
+        const providers = allProviders
           .filter((provider) => !provider.disabled)
           .filter((provider) => provider.weight !== 0)
           .filter((provider) => !retry.excludeProviders.includes(provider.id))
           .filter((provider) => {
             if (!provider.tpmLimit) return true
             const usage = modelTpmLimits?.[`${provider.id}/${provider.model}`] ?? 0
-            return usage < provider.tpmLimit * 1_000_000
+            return usage < provider.tpmLimit * 1_000_000 * 0.8
+          })
+          .map((provider) => {
+            topPriority = Math.min(topPriority, provider.priority)
+            return provider
           })
-
-        const topPriority = Math.min(...allProviders.map((p) => p.priority))
-        const providers = allProviders
           .filter((p) => p.priority <= topPriority)
           .flatMap((provider) => Array<typeof provider>(provider.weight).fill(provider))
 
diff --git a/packages/console/app/src/routes/zen/util/modelTpmLimiter.ts b/packages/console/app/src/routes/zen/util/modelTpmLimiter.ts
index 53015d51c..8e3e8cc95 100644
--- a/packages/console/app/src/routes/zen/util/modelTpmLimiter.ts
+++ b/packages/console/app/src/routes/zen/util/modelTpmLimiter.ts
@@ -1,28 +1,25 @@
 import { and, Database, eq, inArray, sql } from "@opencode-ai/console-core/drizzle/index.js"
-import { ModelTpmLimitTable } from "@opencode-ai/console-core/schema/ip.sql.js"
+import { ModelTpmRateLimitTable } 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 ids = providers.filter((p) => p.tpmLimit).map((p) => `${p.id}/${p.model}`)
   if (ids.length === 0) return
 
-  const yyyyMMddHHmm = new Date(Date.now())
-    .toISOString()
-    .replace(/[^0-9]/g, "")
-    .substring(0, 12)
+  const yyyyMMddHHmm = parseInt(
+    new Date(Date.now())
+      .toISOString()
+      .replace(/[^0-9]/g, "")
+      .substring(0, 12),
+  )
 
   return {
     check: async () => {
       const data = await Database.use((tx) =>
         tx
           .select()
-          .from(ModelTpmLimitTable)
-          .where(
-            inArray(
-              ModelTpmLimitTable.id,
-              ids.map((id) => formatId(id, yyyyMMddHHmm)),
-            ),
-          ),
+          .from(ModelTpmRateLimitTable)
+          .where(and(inArray(ModelTpmRateLimitTable.id, ids), eq(ModelTpmRateLimitTable.interval, yyyyMMddHHmm))),
       )
 
       // convert to map of model to count
@@ -41,14 +38,10 @@ export function createModelTpmLimiter(providers: { id: string; model: string; tp
       if (usage <= 0) return
       await Database.use((tx) =>
         tx
-          .insert(ModelTpmLimitTable)
-          .values({ id: formatId(id, yyyyMMddHHmm), count: usage })
-          .onDuplicateKeyUpdate({ set: { count: sql`${ModelTpmLimitTable.count} + ${usage}` } }),
+          .insert(ModelTpmRateLimitTable)
+          .values({ id, interval: yyyyMMddHHmm, count: usage })
+          .onDuplicateKeyUpdate({ set: { count: sql`${ModelTpmRateLimitTable.count} + ${usage}` } }),
       )
     },
   }
-
-  function formatId(id: string, yyyyMMddHHmm: string) {
-    return `${id.substring(0, 200)}/${yyyyMMddHHmm}`
-  }
 }