zen: use balance after rate limited

3 files changed, 122 insertions, 1 deletions

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
index 766cff684..f9dc0cb45 100644
--- 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
index 03c75387a..6d18f0a2f 100644
--- 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
index b5a90f749..4892ab418 100644
--- 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,