diff options
Diffstat (limited to 'packages/console/app/src')
| -rw-r--r-- | packages/console/app/src/routes/zen/util/handler.ts | 22 |
1 files changed, 13 insertions, 9 deletions
diff --git a/packages/console/app/src/routes/zen/util/handler.ts b/packages/console/app/src/routes/zen/util/handler.ts index 70d2f5516..4346060de 100644 --- a/packages/console/app/src/routes/zen/util/handler.ts +++ b/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()`)), ), ), |