wip: zen black

2 files changed, 29 insertions, 11 deletions

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) => [