wip: zen black

6 files changed, 52 insertions, 22 deletions

diff --git a/packages/console/core/script/lookup-user.ts b/packages/console/core/script/lookup-user.ts
index 3dc5e7a96..355716d1d 100644
--- a/packages/console/core/script/lookup-user.ts
+++ b/packages/console/core/script/lookup-user.ts
@@ -1,7 +1,7 @@
 import { Database, and, eq, sql } from "../src/drizzle/index.js"
 import { AuthTable } from "../src/schema/auth.sql.js"
 import { UserTable } from "../src/schema/user.sql.js"
-import { BillingTable, PaymentTable, SubscriptionTable, UsageTable } from "../src/schema/billing.sql.js"
+import { BillingTable, PaymentTable, SubscriptionTable, SubscriptionPlan, UsageTable } from "../src/schema/billing.sql.js"
 import { WorkspaceTable } from "../src/schema/workspace.sql.js"
 import { BlackData } from "../src/black.js"
 import { centsToMicroCents } from "../src/util/price.js"
@@ -86,8 +86,10 @@ async function printWorkspace(workspaceID: string) {
         timeFixedUpdated: SubscriptionTable.timeFixedUpdated,
         timeRollingUpdated: SubscriptionTable.timeRollingUpdated,
         timeSubscriptionCreated: SubscriptionTable.timeCreated,
+        subscription: BillingTable.subscription,
       })
       .from(UserTable)
+      .innerJoin(BillingTable, eq(BillingTable.workspaceID, workspace.id))
       .leftJoin(AuthTable, and(eq(UserTable.accountID, AuthTable.accountID), eq(AuthTable.provider, "email")))
       .leftJoin(SubscriptionTable, eq(SubscriptionTable.userID, UserTable.id))
       .where(eq(UserTable.workspaceID, workspace.id))
@@ -223,17 +225,20 @@ function formatRetryTime(seconds: number) {
 }
 
 function getSubscriptionStatus(row: {
+  subscription: {
+    plan: typeof SubscriptionPlan[number]
+  } | null
   timeSubscriptionCreated: Date | null
   fixedUsage: number | null
   rollingUsage: number | null
   timeFixedUpdated: Date | null
   timeRollingUpdated: Date | null
 }) {
-  if (!row.timeSubscriptionCreated) {
+  if (!row.timeSubscriptionCreated || !row.subscription) {
     return { weekly: null, rolling: null, rateLimited: null, retryIn: null }
   }
 
-  const black = BlackData.get()
+  const black = BlackData.get({ plan: row.subscription.plan })
   const now = new Date()
   const week = getWeekBounds(now)
 
diff --git a/packages/console/core/script/promote-black.ts b/packages/console/core/script/promote-black.ts
index bb3dcc6f7..4338d0e42 100755
--- a/packages/console/core/script/promote-black.ts
+++ b/packages/console/core/script/promote-black.ts
@@ -12,11 +12,11 @@ const root = path.resolve(process.cwd(), "..", "..", "..")
 // read the secret
 const ret = await $`bun sst secret list`.cwd(root).text()
 const lines = ret.split("\n")
-const value = lines.find((line) => line.startsWith("ZEN_BLACK"))?.split("=")[1]
-if (!value) throw new Error("ZEN_BLACK not found")
+const value = lines.find((line) => line.startsWith("ZEN_BLACK_LIMITS"))?.split("=")[1]
+if (!value) throw new Error("ZEN_BLACK_LIMITS not found")
 
 // validate value
 BlackData.validate(JSON.parse(value))
 
 // update the secret
-await $`bun sst secret set ZEN_BLACK ${value} --stage ${stage}`
+await $`bun sst secret set ZEN_BLACK_LIMITS ${value} --stage ${stage}`
diff --git a/packages/console/core/script/update-black.ts b/packages/console/core/script/update-black.ts
index 58923b457..695a5d3ce 100755
--- a/packages/console/core/script/update-black.ts
+++ b/packages/console/core/script/update-black.ts
@@ -8,10 +8,10 @@ import { BlackData } from "../src/black"
 const root = path.resolve(process.cwd(), "..", "..", "..")
 const secrets = await $`bun sst secret list`.cwd(root).text()
 
-// read the line starting with "ZEN_BLACK"
+// read value
 const lines = secrets.split("\n")
-const oldValue = lines.find((line) => line.startsWith("ZEN_BLACK"))?.split("=")[1]
-if (!oldValue) throw new Error("ZEN_BLACK not found")
+const oldValue = lines.find((line) => line.startsWith("ZEN_BLACK_LIMITS"))?.split("=")[1] ?? "{}"
+if (!oldValue) throw new Error("ZEN_BLACK_LIMITS not found")
 
 // store the prettified json to a temp file
 const filename = `black-${Date.now()}.json`
@@ -25,4 +25,4 @@ const newValue = JSON.stringify(JSON.parse(await tempFile.text()))
 BlackData.validate(JSON.parse(newValue))
 
 // update the secret
-await $`bun sst secret set ZEN_BLACK ${newValue}`
+await $`bun sst secret set ZEN_BLACK_LIMITS ${newValue}`
diff --git a/packages/console/core/src/black.ts b/packages/console/core/src/black.ts
index 753d25808..85e16f1a1 100644
--- a/packages/console/core/src/black.ts
+++ b/packages/console/core/src/black.ts
@@ -3,33 +3,49 @@ import { fn } from "./util/fn"
 import { Resource } from "@opencode-ai/console-resource"
 import { centsToMicroCents } from "./util/price"
 import { getWeekBounds } from "./util/date"
+import { SubscriptionPlan } from "./schema/billing.sql"
 
 export namespace BlackData {
   const Schema = z.object({
-    fixedLimit: z.number().int(),
-    rollingLimit: z.number().int(),
-    rollingWindow: z.number().int(),
+    "200": z.object({
+      fixedLimit: z.number().int(),
+      rollingLimit: z.number().int(),
+      rollingWindow: z.number().int(),
+    }),
+    "100": z.object({
+      fixedLimit: z.number().int(),
+      rollingLimit: z.number().int(),
+      rollingWindow: z.number().int(),
+    }),
+    "20": z.object({
+      fixedLimit: z.number().int(),
+      rollingLimit: z.number().int(),
+      rollingWindow: z.number().int(),
+    }),
   })
 
   export const validate = fn(Schema, (input) => {
     return input
   })
 
-  export const get = fn(z.void(), () => {
-    const json = JSON.parse(Resource.ZEN_BLACK.value)
-    return Schema.parse(json)
+  export const get = fn(z.object({
+      plan: z.enum(SubscriptionPlan),
+    }), ({ plan }) => {
+    const json = JSON.parse(Resource.ZEN_BLACK_LIMITS.value)
+    return Schema.parse(json)[plan]
   })
 }
 
 export namespace Black {
   export const analyzeRollingUsage = fn(
     z.object({
+      plan: z.enum(SubscriptionPlan),
       usage: z.number().int(),
       timeUpdated: z.date(),
     }),
-    ({ usage, timeUpdated }) => {
+    ({ plan, usage, timeUpdated }) => {
       const now = new Date()
-      const black = BlackData.get()
+      const black = BlackData.get({ plan })
       const rollingWindowMs = black.rollingWindow * 3600 * 1000
       const rollingLimitInMicroCents = centsToMicroCents(black.rollingLimit * 100)
       const windowStart = new Date(now.getTime() - rollingWindowMs)
@@ -59,11 +75,12 @@ export namespace Black {
 
   export const analyzeWeeklyUsage = fn(
     z.object({
+      plan:z.enum(SubscriptionPlan),
       usage: z.number().int(),
       timeUpdated: z.date(),
     }),
-    ({ usage, timeUpdated }) => {
-      const black = BlackData.get()
+    ({ plan, usage, timeUpdated }) => {
+      const black = BlackData.get({ plan })
       const now = new Date()
       const week = getWeekBounds(now)
       const fixedLimitInMicroCents = centsToMicroCents(black.fixedLimit * 100)
diff --git a/packages/console/core/src/schema/billing.sql.ts b/packages/console/core/src/schema/billing.sql.ts
index 9f05919f2..ae32ed5ce 100644
--- a/packages/console/core/src/schema/billing.sql.ts
+++ b/packages/console/core/src/schema/billing.sql.ts
@@ -2,6 +2,7 @@ import { bigint, boolean, index, int, json, mysqlEnum, mysqlTable, uniqueIndex,
 import { timestamps, ulid, utc, workspaceColumns } from "../drizzle/types"
 import { workspaceIndexes } from "./workspace.sql"
 
+export const SubscriptionPlan = ["20", "100", "200"] as const
 export const BillingTable = mysqlTable(
   "billing",
   {
@@ -28,7 +29,7 @@ export const BillingTable = mysqlTable(
       plan: "20" | "100" | "200"
     }>(),
     subscriptionID: varchar("subscription_id", { length: 28 }),
-    subscriptionPlan: mysqlEnum("subscription_plan", ["20", "100", "200"] as const),
+    subscriptionPlan: mysqlEnum("subscription_plan", SubscriptionPlan),
     timeSubscriptionBooked: utc("time_subscription_booked"),
   },
   (table) => [
diff --git a/packages/console/core/sst-env.d.ts b/packages/console/core/sst-env.d.ts
index 66c06a824..cb4902813 100644
--- a/packages/console/core/sst-env.d.ts
+++ b/packages/console/core/sst-env.d.ts
@@ -118,10 +118,17 @@ declare module "sst" {
       "type": "sst.cloudflare.StaticSite"
       "url": string
     }
-    "ZEN_BLACK": {
+    "ZEN_BLACK_LIMITS": {
       "type": "sst.sst.Secret"
       "value": string
     }
+    "ZEN_BLACK_PRICE": {
+      "plan100": string
+      "plan20": string
+      "plan200": string
+      "product": string
+      "type": "sst.sst.Linkable"
+    }
     "ZEN_MODELS1": {
       "type": "sst.sst.Secret"
       "value": string