zen: use balance after rate limited

2026-01-23 10:57:58 -05:00
parent 65c236c071
commit 24d942349f
5 changed files with 181 additions and 51 deletions
--- 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;
+    }
+  }
 }
--- 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>
--- 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,
--- 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
+      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`
+        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 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 "subscription"
+      } catch(e) {
+        if (!authInfo.billing.subscription.useBalance) throw e
      }
-
-      // 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,
-          )
-      }
-
-      return
    }

    // 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 })
--- 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),