zen: optimize query

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

@@ -172,8 +172,8 @@ export async function handler(
       const tokensInfo = providerInfo.normalizeUsage(json.usage)
       await trialLimiter?.track(tokensInfo)
       await rateLimiter?.track()
-      await trackUsage(authInfo, modelInfo, providerInfo, tokensInfo)
-      await reload(authInfo)
+      const costInfo = await trackUsage(authInfo, modelInfo, providerInfo, tokensInfo)
+      await reload(authInfo, costInfo)
       return new Response(body, {
         status: resStatus,
         statusText: res.statusText,
@@ -206,8 +206,8 @@ export async function handler(
                 if (usage) {
                   const tokensInfo = providerInfo.normalizeUsage(usage)
                   await trialLimiter?.track(tokensInfo)
-                  await trackUsage(authInfo, modelInfo, providerInfo, tokensInfo)
-                  await reload(authInfo)
+                  const costInfo = await trackUsage(authInfo, modelInfo, providerInfo, tokensInfo)
+                  await reload(authInfo, costInfo)
                 }
                 c.close()
                 return
@@ -608,13 +608,20 @@ export async function handler(
         .set({ timeUsed: sql`now()` })
         .where(and(eq(KeyTable.workspaceID, authInfo.workspaceID), eq(KeyTable.id, authInfo.apiKeyId))),
     )
+
+    return { costInMicroCents: cost }
   }
 
-  async function reload(authInfo: AuthInfo) {
+  async function reload(authInfo: AuthInfo, costInfo: Awaited<ReturnType<typeof trackUsage>>) {
     if (!authInfo) return
     if (authInfo.isFree) return
     if (authInfo.provider?.credentials) return
 
+    if (!costInfo) return
+
+    const reloadTrigger = centsToMicroCents((authInfo.billing.reloadTrigger ?? Billing.RELOAD_TRIGGER) * 100)
+    if (authInfo.billing.balance - costInfo.costInMicroCents >= reloadTrigger) return
+
     const lock = await Database.use((tx) =>
       tx
         .update(BillingTable)
@@ -625,10 +632,7 @@ export async function handler(
           and(
             eq(BillingTable.workspaceID, authInfo.workspaceID),
             eq(BillingTable.reload, true),
-            lt(
-              BillingTable.balance,
-              centsToMicroCents((authInfo.billing.reloadTrigger ?? Billing.RELOAD_TRIGGER) * 100),
-            ),
+            lt(BillingTable.balance, reloadTrigger),
             or(isNull(BillingTable.timeReloadLockedTill), lt(BillingTable.timeReloadLockedTill, sql`now()`)),
           ),
         ),