zen: use balance after rate limited

fwang · fwang · commit 24d942349f78 · 2026-01-23T10:58:00.000-05:00
diff --git a/packages/console/app/src/routes/workspace/[id]/billing/black-section.module.css b/packages/console/app/src/routes/workspace/[id]/billing/black-section.module.css
@@ -59,4 +59,84 @@
     font-size: var(--font-size-sm);
     color: var(--color-text-muted);
   }
+
+  [data-slot="setting-row"] {
+    display: flex;
+    align-items: center;
+    justify-content: space-between;
+    gap: var(--space-3);
+    margin-top: var(--space-4);
+
+    p {
+      font-size: var(--font-size-sm);
+      line-height: 1.5;
+      color: var(--color-text-secondary);
+      margin: 0;
+    }
+  }
+
+  [data-slot="toggle-label"] {
+    position: relative;
+    display: inline-block;
+    width: 2.5rem;
+    height: 1.5rem;
+    cursor: pointer;
+    flex-shrink: 0;
+
+    input {
+      opacity: 0;
+      width: 0;
+      height: 0;
+    }
+
+    span {
+      position: absolute;
+      inset: 0;
+      background-color: #ccc;
+      border: 1px solid #bbb;
+      border-radius: 1.5rem;
+      transition: all 0.3s ease;
+      cursor: pointer;
+
+      &::before {
+        content: "";
+        position: absolute;
+        top: 50%;
+        left: 0.125rem;
+        width: 1.25rem;
+        height: 1.25rem;
+        background-color: white;
+        border: 1px solid #ddd;
+        border-radius: 50%;
+        transform: translateY(-50%);
+        transition: all 0.3s ease;
+      }
+    }
+
+    input:checked + span {
+      background-color: #21ad0e;
+      border-color: #148605;
+
+      &::before {
+        transform: translateX(1rem) translateY(-50%);
+      }
+    }
+
+    &:hover span {
+      box-shadow: 0 0 0 2px rgba(34, 197, 94, 0.2);
+    }
+
+    input:checked:hover + span {
+      box-shadow: 0 0 0 2px rgba(34, 197, 94, 0.3);
+    }
+
+    &:has(input:disabled) {
+      cursor: not-allowed;
+    }
+
+    input:disabled + span {
+      opacity: 0.5;
+      cursor: not-allowed;
+    }
+  }
 }
diff --git a/packages/console/app/src/routes/workspace/[id]/billing/black-section.tsx b/packages/console/app/src/routes/workspace/[id]/billing/black-section.tsx
@@ -2,7 +2,7 @@ import { action, useParams, useAction, useSubmission, json, query, createAsync }
 import { createStore } from "solid-js/store"
 import { Show } from "solid-js"
 import { Billing } from "@opencode-ai/console-core/billing.js"
-import { Database, eq, and, isNull } from "@opencode-ai/console-core/drizzle/index.js"
+import { Database, eq, and, isNull, sql } from "@opencode-ai/console-core/drizzle/index.js"
 import { BillingTable, SubscriptionTable } from "@opencode-ai/console-core/schema/billing.sql.js"
 import { Actor } from "@opencode-ai/console-core/actor.js"
 import { Black } from "@opencode-ai/console-core/black.js"
@@ -32,6 +32,7 @@ const querySubscription = query(async (workspaceID: string) => {
 
     return {
       plan: row.subscription.plan,
+      useBalance: row.subscription.useBalance ?? false,
       rollingUsage: Black.analyzeRollingUsage({
         plan: row.subscription.plan,
         usage: row.rollingUsage ?? 0,
@@ -107,6 +108,30 @@ const createSessionUrl = action(async (workspaceID: string, returnUrl: string) =
   )
 }, "sessionUrl")
 
+const setUseBalance = action(async (form: FormData) => {
+  "use server"
+  const workspaceID = form.get("workspaceID")?.toString()
+  if (!workspaceID) return { error: "Workspace ID is required" }
+  const useBalance = form.get("useBalance")?.toString() === "true"
+
+  return json(
+    await withActor(async () => {
+      await Database.use((tx) =>
+        tx
+          .update(BillingTable)
+          .set({
+            subscription: useBalance
+              ? sql`JSON_SET(subscription, '$.useBalance', true)`
+              : sql`JSON_REMOVE(subscription, '$.useBalance')`,
+          })
+          .where(eq(BillingTable.workspaceID, workspaceID)),
+      )
+      return { error: undefined }
+    }, workspaceID).catch((e) => ({ error: e.message as string })),
+    { revalidate: [queryBillingInfo.key, querySubscription.key] },
+  )
+}, "setUseBalance")
+
 export function BlackSection() {
   const params = useParams()
   const billing = createAsync(() => queryBillingInfo(params.id!))
@@ -117,6 +142,7 @@ export function BlackSection() {
   const cancelSubmission = useSubmission(cancelWaitlist)
   const enrollAction = useAction(enroll)
   const enrollSubmission = useSubmission(enroll)
+  const useBalanceSubmission = useSubmission(setUseBalance)
   const [store, setStore] = createStore({
     sessionRedirecting: false,
     cancelled: false,
@@ -185,6 +211,20 @@ export function BlackSection() {
                 <span data-slot="reset-time">Resets in {formatResetTime(sub().weeklyUsage.resetInSec)}</span>
               </div>
             </div>
+            <form action={setUseBalance} method="post" data-slot="setting-row">
+              <p>Use your available balance after reaching the usage limits</p>
+              <input type="hidden" name="workspaceID" value={params.id} />
+              <input type="hidden" name="useBalance" value={sub().useBalance ? "false" : "true"} />
+              <label data-slot="toggle-label">
+                <input
+                  type="checkbox"
+                  checked={sub().useBalance}
+                  disabled={useBalanceSubmission.pending}
+                  onChange={(e) => e.currentTarget.form?.requestSubmit()}
+                />
+                <span></span>
+              </label>
+            </form>
           </section>
         )}
       </Show>
diff --git a/packages/console/app/src/routes/workspace/common.tsx b/packages/console/app/src/routes/workspace/common.tsx
@@ -110,6 +110,7 @@ export const queryBillingInfo = query(async (workspaceID: string) => {
       timeMonthlyUsageUpdated: billing.timeMonthlyUsageUpdated,
       reloadError: billing.reloadError,
       timeReloadError: billing.timeReloadError,
+      subscription: billing.subscription,
       subscriptionID: billing.subscriptionID,
       subscriptionPlan: billing.subscriptionPlan,
       timeSubscriptionBooked: billing.timeSubscriptionBooked,
diff --git a/packages/console/app/src/routes/zen/util/handler.ts b/packages/console/app/src/routes/zen/util/handler.ts
@@ -84,6 +84,7 @@ export async function handler(
     const stickyTracker = createStickyTracker(modelInfo.stickyProvider, sessionId)
     const stickyProvider = await stickyTracker?.get()
     const authInfo = await authenticate(modelInfo)
+    const billingSource = validateBilling(authInfo, modelInfo)
 
     const retriableRequest = async (retry: RetryOptions = { excludeProviders: [], retryCount: 0 }) => {
       const providerInfo = selectProvider(
@@ -96,7 +97,6 @@ export async function handler(
         retry,
         stickyProvider,
       )
-      validateBilling(authInfo, modelInfo)
       validateModelSettings(authInfo)
       updateProviderKey(authInfo, providerInfo)
       logger.metric({ provider: providerInfo.id })
@@ -183,7 +183,7 @@ export async function handler(
       const tokensInfo = providerInfo.normalizeUsage(json.usage)
       await trialLimiter?.track(tokensInfo)
       await rateLimiter?.track()
-      const costInfo = await trackUsage(authInfo, modelInfo, providerInfo, tokensInfo)
+      const costInfo = await trackUsage(authInfo, modelInfo, providerInfo, billingSource, tokensInfo)
       await reload(authInfo, costInfo)
       return new Response(body, {
         status: resStatus,
@@ -219,7 +219,7 @@ export async function handler(
                 if (usage) {
                   const tokensInfo = providerInfo.normalizeUsage(usage)
                   await trialLimiter?.track(tokensInfo)
-                  const costInfo = await trackUsage(authInfo, modelInfo, providerInfo, tokensInfo)
+                  const costInfo = await trackUsage(authInfo, modelInfo, providerInfo, billingSource, tokensInfo)
                   await reload(authInfo, costInfo)
                 }
                 c.close()
@@ -484,54 +484,58 @@ export async function handler(
   }
 
   function validateBilling(authInfo: AuthInfo, modelInfo: ModelInfo) {
-    if (!authInfo) return
-    if (authInfo.provider?.credentials) return
-    if (authInfo.isFree) return
-    if (modelInfo.allowAnonymous) return
+    if (!authInfo) return "anonymous"
+    if (authInfo.provider?.credentials) return "free"
+    if (authInfo.isFree) return "free"
+    if (modelInfo.allowAnonymous) return "free"
 
     // Validate subscription billing
     if (authInfo.billing.subscription && authInfo.subscription) {
-      const sub = authInfo.subscription
-      const plan = authInfo.billing.subscription.plan
-
-      const formatRetryTime = (seconds: number) => {
-        const days = Math.floor(seconds / 86400)
-        if (days >= 1) return `${days} day${days > 1 ? "s" : ""}`
-        const hours = Math.floor(seconds / 3600)
-        const minutes = Math.ceil((seconds % 3600) / 60)
-        if (hours >= 1) return `${hours}hr ${minutes}min`
-        return `${minutes}min`
-      }
+      try {
+        const sub = authInfo.subscription
+        const plan = authInfo.billing.subscription.plan
+
+        const formatRetryTime = (seconds: number) => {
+          const days = Math.floor(seconds / 86400)
+          if (days >= 1) return `${days} day${days > 1 ? "s" : ""}`
+          const hours = Math.floor(seconds / 3600)
+          const minutes = Math.ceil((seconds % 3600) / 60)
+          if (hours >= 1) return `${hours}hr ${minutes}min`
+          return `${minutes}min`
+        }
 
-      // Check weekly limit
-      if (sub.fixedUsage && sub.timeFixedUpdated) {
-        const result = Black.analyzeWeeklyUsage({
-          plan,
-          usage: sub.fixedUsage,
-          timeUpdated: sub.timeFixedUpdated,
-        })
-        if (result.status === "rate-limited")
-          throw new SubscriptionError(
-            `Subscription quota exceeded. Retry in ${formatRetryTime(result.resetInSec)}.`,
-            result.resetInSec,
-          )
-      }
+        // Check weekly limit
+        if (sub.fixedUsage && sub.timeFixedUpdated) {
+          const result = Black.analyzeWeeklyUsage({
+            plan,
+            usage: sub.fixedUsage,
+            timeUpdated: sub.timeFixedUpdated,
+          })
+          if (result.status === "rate-limited")
+            throw new SubscriptionError(
+              `Subscription quota exceeded. Retry in ${formatRetryTime(result.resetInSec)}.`,
+              result.resetInSec,
+            )
+        }
 
-      // Check rolling limit
-      if (sub.rollingUsage && sub.timeRollingUpdated) {
-        const result = Black.analyzeRollingUsage({
-          plan,
-          usage: sub.rollingUsage,
-          timeUpdated: sub.timeRollingUpdated,
-        })
-        if (result.status === "rate-limited")
-          throw new SubscriptionError(
-            `Subscription quota exceeded. Retry in ${formatRetryTime(result.resetInSec)}.`,
-            result.resetInSec,
-          )
-      }
+        // Check rolling limit
+        if (sub.rollingUsage && sub.timeRollingUpdated) {
+          const result = Black.analyzeRollingUsage({
+            plan,
+            usage: sub.rollingUsage,
+            timeUpdated: sub.timeRollingUpdated,
+          })
+          if (result.status === "rate-limited")
+            throw new SubscriptionError(
+              `Subscription quota exceeded. Retry in ${formatRetryTime(result.resetInSec)}.`,
+              result.resetInSec,
+            )
+        }
 
-      return
+        return "subscription"
+      } catch(e) {
+        if (!authInfo.billing.subscription.useBalance) throw e
+      }
     }
 
     // Validate pay as you go billing
@@ -571,6 +575,8 @@ export async function handler(
       throw new UserLimitError(
         `You have reached your monthly spending limit of $${authInfo.user.monthlyLimit}. Manage your limits here: https://opencode.ai/workspace/${authInfo.workspaceID}/members`,
       )
+
+    return "balance"
   }
 
   function validateModelSettings(authInfo: AuthInfo) {
@@ -587,6 +593,7 @@ export async function handler(
     authInfo: AuthInfo,
     modelInfo: ModelInfo,
     providerInfo: ProviderInfo,
+    billingSource: ReturnType<typeof validateBilling>,
     usageInfo: UsageInfo,
   ) {
     const { inputTokens, outputTokens, reasoningTokens, cacheReadTokens, cacheWrite5mTokens, cacheWrite1hTokens } =
@@ -643,7 +650,8 @@ export async function handler(
       "cost.total": Math.round(totalCostInCent),
     })
 
-    if (!authInfo) return
+    if (billingSource === "anonymous") return
+    authInfo = authInfo!
 
     const cost = authInfo.provider?.credentials ? 0 : centsToMicroCents(totalCostInCent)
     await Database.use((db) =>
@@ -661,13 +669,13 @@ export async function handler(
           cacheWrite1hTokens,
           cost,
           keyID: authInfo.apiKeyId,
-          enrichment: authInfo.subscription ? { plan: "sub" } : undefined,
+          enrichment: billingSource === "subscription" ? { plan: "sub" } : undefined,
         }),
         db
           .update(KeyTable)
           .set({ timeUsed: sql`now()` })
           .where(and(eq(KeyTable.workspaceID, authInfo.workspaceID), eq(KeyTable.id, authInfo.apiKeyId))),
-        ...(authInfo.subscription
+        ...(billingSource === "subscription"
           ? (() => {
               const plan = authInfo.billing.subscription!.plan
               const black = BlackData.getLimits({ plan })
diff --git a/packages/console/core/src/schema/billing.sql.ts b/packages/console/core/src/schema/billing.sql.ts
@@ -24,9 +24,10 @@ export const BillingTable = mysqlTable(
     timeReloadLockedTill: utc("time_reload_locked_till"),
     subscription: json("subscription").$type<{
       status: "subscribed"
-      coupon?: string
       seats: number
       plan: "20" | "100" | "200"
+      useBalance?: boolean
+      coupon?: string
     }>(),
     subscriptionID: varchar("subscription_id", { length: 28 }),
     subscriptionPlan: mysqlEnum("subscription_plan", SubscriptionPlan),
