wip: zen

5 files changed, 323 insertions, 61 deletions

diff --git a/packages/console/core/script/black-stats.ts b/packages/console/core/script/black-stats.ts
new file mode 100644
index 000000000..de7cf5e41
--- /dev/null
+++ b/packages/console/core/script/black-stats.ts
@@ -0,0 +1,312 @@
+import { Database, and, eq, inArray, isNotNull, sql } from "../src/drizzle/index.js"
+import { BillingTable, BlackPlans, SubscriptionTable, UsageTable } from "../src/schema/billing.sql.js"
+
+if (process.argv.length < 3) {
+  console.error("Usage: bun black-stats.ts <plan>")
+  process.exit(1)
+}
+const plan = process.argv[2] as (typeof BlackPlans)[number]
+if (!BlackPlans.includes(plan)) {
+  console.error("Usage: bun black-stats.ts <plan>")
+  process.exit(1)
+}
+const cutoff = new Date(Date.UTC(2026, 1, 0, 23, 59, 59, 999))
+
+// get workspaces
+const workspaces = await Database.use((tx) =>
+  tx
+    .select({ workspaceID: BillingTable.workspaceID })
+    .from(BillingTable)
+    .where(
+      and(isNotNull(BillingTable.subscriptionID), sql`JSON_UNQUOTE(JSON_EXTRACT(subscription, '$.plan')) = ${plan}`),
+    ),
+)
+if (workspaces.length === 0) throw new Error(`No active Black ${plan} subscriptions found`)
+
+const week = sql<number>`YEARWEEK(${UsageTable.timeCreated}, 3)`
+const workspaceIDs = workspaces.map((row) => row.workspaceID)
+// Get subscription spend
+const spend = await Database.use((tx) =>
+  tx
+    .select({
+      workspaceID: UsageTable.workspaceID,
+      week,
+      amount: sql<number>`COALESCE(SUM(${UsageTable.cost}), 0)`,
+    })
+    .from(UsageTable)
+    .where(
+      and(inArray(UsageTable.workspaceID, workspaceIDs), sql`JSON_UNQUOTE(JSON_EXTRACT(enrichment, '$.plan')) = 'sub'`),
+    )
+    .groupBy(UsageTable.workspaceID, week),
+)
+
+// Get pay per use spend
+const ppu = await Database.use((tx) =>
+  tx
+    .select({
+      workspaceID: UsageTable.workspaceID,
+      week,
+      amount: sql<number>`COALESCE(SUM(${UsageTable.cost}), 0)`,
+    })
+    .from(UsageTable)
+    .where(
+      and(
+        inArray(UsageTable.workspaceID, workspaceIDs),
+        sql`(${UsageTable.enrichment} IS NULL OR JSON_UNQUOTE(JSON_EXTRACT(enrichment, '$.plan')) != 'sub')`,
+      ),
+    )
+    .groupBy(UsageTable.workspaceID, week),
+)
+
+const models = await Database.use((tx) =>
+  tx
+    .select({
+      workspaceID: UsageTable.workspaceID,
+      model: UsageTable.model,
+      amount: sql<number>`COALESCE(SUM(${UsageTable.cost}), 0)`,
+    })
+    .from(UsageTable)
+    .where(
+      and(inArray(UsageTable.workspaceID, workspaceIDs), sql`JSON_UNQUOTE(JSON_EXTRACT(enrichment, '$.plan')) = 'sub'`),
+    )
+    .groupBy(UsageTable.workspaceID, UsageTable.model),
+)
+
+const tokens = await Database.use((tx) =>
+  tx
+    .select({
+      workspaceID: UsageTable.workspaceID,
+      week,
+      input: sql<number>`COALESCE(SUM(${UsageTable.inputTokens}), 0)`,
+      cacheRead: sql<number>`COALESCE(SUM(${UsageTable.cacheReadTokens}), 0)`,
+      output: sql<number>`COALESCE(SUM(${UsageTable.outputTokens}), 0) + COALESCE(SUM(${UsageTable.reasoningTokens}), 0)`,
+    })
+    .from(UsageTable)
+    .where(
+      and(inArray(UsageTable.workspaceID, workspaceIDs), sql`JSON_UNQUOTE(JSON_EXTRACT(enrichment, '$.plan')) = 'sub'`),
+    )
+    .groupBy(UsageTable.workspaceID, week),
+)
+
+const allWeeks = [...spend, ...ppu].map((row) => row.week)
+const weeks = [...new Set(allWeeks)].sort((a, b) => a - b)
+const spendMap = new Map<string, Map<number, number>>()
+const totals = new Map<string, number>()
+const ppuMap = new Map<string, Map<number, number>>()
+const ppuTotals = new Map<string, number>()
+const modelMap = new Map<string, { model: string; amount: number }[]>()
+const tokenMap = new Map<string, Map<number, { input: number; cacheRead: number; output: number }>>()
+
+for (const row of spend) {
+  const workspace = spendMap.get(row.workspaceID) ?? new Map<number, number>()
+  const total = totals.get(row.workspaceID) ?? 0
+  const amount = toNumber(row.amount)
+  workspace.set(row.week, amount)
+  totals.set(row.workspaceID, total + amount)
+  spendMap.set(row.workspaceID, workspace)
+}
+
+for (const row of ppu) {
+  const workspace = ppuMap.get(row.workspaceID) ?? new Map<number, number>()
+  const total = ppuTotals.get(row.workspaceID) ?? 0
+  const amount = toNumber(row.amount)
+  workspace.set(row.week, amount)
+  ppuTotals.set(row.workspaceID, total + amount)
+  ppuMap.set(row.workspaceID, workspace)
+}
+
+for (const row of models) {
+  const current = modelMap.get(row.workspaceID) ?? []
+  current.push({ model: row.model, amount: toNumber(row.amount) })
+  modelMap.set(row.workspaceID, current)
+}
+
+for (const row of tokens) {
+  const workspace = tokenMap.get(row.workspaceID) ?? new Map()
+  workspace.set(row.week, {
+    input: toNumber(row.input),
+    cacheRead: toNumber(row.cacheRead),
+    output: toNumber(row.output),
+  })
+  tokenMap.set(row.workspaceID, workspace)
+}
+
+const users = await Database.use((tx) =>
+  tx
+    .select({
+      workspaceID: SubscriptionTable.workspaceID,
+      subscribed: SubscriptionTable.timeCreated,
+      subscription: BillingTable.subscription,
+    })
+    .from(SubscriptionTable)
+    .innerJoin(BillingTable, eq(SubscriptionTable.workspaceID, BillingTable.workspaceID))
+    .where(
+      and(inArray(SubscriptionTable.workspaceID, workspaceIDs), sql`${SubscriptionTable.timeCreated} <= ${cutoff}`),
+    ),
+)
+
+const counts = new Map<string, number>()
+for (const user of users) {
+  const current = counts.get(user.workspaceID) ?? 0
+  counts.set(user.workspaceID, current + 1)
+}
+
+const rows = users
+  .map((user) => {
+    const workspace = spendMap.get(user.workspaceID) ?? new Map<number, number>()
+    const ppuWorkspace = ppuMap.get(user.workspaceID) ?? new Map<number, number>()
+    const count = counts.get(user.workspaceID) ?? 1
+    const amount = (totals.get(user.workspaceID) ?? 0) / count
+    const ppuAmount = (ppuTotals.get(user.workspaceID) ?? 0) / count
+    const monthStart = user.subscribed ? startOfMonth(user.subscribed) : null
+    const modelRows = (modelMap.get(user.workspaceID) ?? []).sort((a, b) => b.amount - a.amount).slice(0, 3)
+    const modelTotal = totals.get(user.workspaceID) ?? 0
+    const modelCells = modelRows.map((row) => ({
+      model: row.model,
+      percent: modelTotal > 0 ? `${((row.amount / modelTotal) * 100).toFixed(1)}%` : "0.0%",
+    }))
+    const modelData = [0, 1, 2].map((index) => modelCells[index] ?? { model: "-", percent: "-" })
+    const weekly = Object.fromEntries(
+      weeks.map((item) => {
+        const value = (workspace.get(item) ?? 0) / count
+        const beforeMonth = monthStart ? isoWeekStart(item) < monthStart : false
+        return [formatWeek(item), beforeMonth ? "-" : formatMicroCents(value)]
+      }),
+    )
+    const ppuWeekly = Object.fromEntries(
+      weeks.map((item) => {
+        const value = (ppuWorkspace.get(item) ?? 0) / count
+        const beforeMonth = monthStart ? isoWeekStart(item) < monthStart : false
+        return [formatWeek(item), beforeMonth ? "-" : formatMicroCents(value)]
+      }),
+    )
+    const tokenWorkspace = tokenMap.get(user.workspaceID) ?? new Map()
+    const weeklyTokens = Object.fromEntries(
+      weeks.map((item) => {
+        const t = tokenWorkspace.get(item) ?? { input: 0, cacheRead: 0, output: 0 }
+        const beforeMonth = monthStart ? isoWeekStart(item) < monthStart : false
+        return [
+          formatWeek(item),
+          beforeMonth
+            ? { input: "-", cacheRead: "-", output: "-" }
+            : {
+                input: Math.round(t.input / count),
+                cacheRead: Math.round(t.cacheRead / count),
+                output: Math.round(t.output / count),
+              },
+        ]
+      }),
+    )
+    return {
+      workspaceID: user.workspaceID,
+      useBalance: user.subscription?.useBalance ?? false,
+      subscribed: formatDate(user.subscribed),
+      subscribedAt: user.subscribed?.getTime() ?? 0,
+      amount,
+      ppuAmount,
+      models: modelData,
+      weekly,
+      ppuWeekly,
+      weeklyTokens,
+    }
+  })
+  .sort((a, b) => a.subscribedAt - b.subscribedAt)
+
+console.log(`Black ${plan} subscribers: ${rows.length}`)
+const header = [
+  "workspaceID",
+  "subscribed",
+  "useCredit",
+  "subTotal",
+  "ppuTotal",
+  "model1",
+  "model1%",
+  "model2",
+  "model2%",
+  "model3",
+  "model3%",
+  ...weeks.flatMap((item) => [
+    formatWeek(item) + " sub",
+    formatWeek(item) + " ppu",
+    formatWeek(item) + " input",
+    formatWeek(item) + " cache",
+    formatWeek(item) + " output",
+  ]),
+]
+const lines = [header.map(csvCell).join(",")]
+for (const row of rows) {
+  const model1 = row.models[0]
+  const model2 = row.models[1]
+  const model3 = row.models[2]
+  const cells = [
+    row.workspaceID,
+    row.subscribed ?? "",
+    row.useBalance ? "yes" : "no",
+    formatMicroCents(row.amount),
+    formatMicroCents(row.ppuAmount),
+    model1.model,
+    model1.percent,
+    model2.model,
+    model2.percent,
+    model3.model,
+    model3.percent,
+    ...weeks.flatMap((item) => {
+      const t = row.weeklyTokens[formatWeek(item)] ?? { input: "-", cacheRead: "-", output: "-" }
+      return [
+        row.weekly[formatWeek(item)] ?? "",
+        row.ppuWeekly[formatWeek(item)] ?? "",
+        String(t.input),
+        String(t.cacheRead),
+        String(t.output),
+      ]
+    }),
+  ]
+  lines.push(cells.map(csvCell).join(","))
+}
+const output = `${lines.join("\n")}\n`
+const file = Bun.file(`black-stats-${plan}.csv`)
+await file.write(output)
+console.log(`Wrote ${lines.length - 1} rows to ${file.name}`)
+const total = rows.reduce((sum, row) => sum + row.amount, 0)
+const average = rows.length === 0 ? 0 : total / rows.length
+console.log(`Average spending per user: ${formatMicroCents(average)}`)
+
+function formatMicroCents(value: number) {
+  return `$${(value / 100000000).toFixed(2)}`
+}
+
+function formatDate(value: Date | null | undefined) {
+  if (!value) return null
+  return value.toISOString().split("T")[0]
+}
+
+function formatWeek(value: number) {
+  return formatDate(isoWeekStart(value)) ?? ""
+}
+
+function startOfMonth(value: Date) {
+  return new Date(Date.UTC(value.getUTCFullYear(), value.getUTCMonth(), 1))
+}
+
+function isoWeekStart(value: number) {
+  const year = Math.floor(value / 100)
+  const weekNumber = value % 100
+  const jan4 = new Date(Date.UTC(year, 0, 4))
+  const day = jan4.getUTCDay() || 7
+  const weekStart = new Date(Date.UTC(year, 0, 4 - (day - 1)))
+  weekStart.setUTCDate(weekStart.getUTCDate() + (weekNumber - 1) * 7)
+  return weekStart
+}
+
+function toNumber(value: unknown) {
+  if (typeof value === "number") return value
+  if (typeof value === "bigint") return Number(value)
+  if (typeof value === "string") return Number(value)
+  return 0
+}
+
+function csvCell(value: string | number) {
+  const text = String(value)
+  if (!/[",\n]/.test(text)) return text
+  return `"${text.replace(/"/g, '""')}"`
+}
diff --git a/packages/console/core/script/promote-black.ts b/packages/console/core/script/promote-black.ts
deleted file mode 100755
index 4338d0e42..000000000
--- a/packages/console/core/script/promote-black.ts
+++ /dev/null
@@ -1,22 +0,0 @@
-#!/usr/bin/env bun
-
-import { $ } from "bun"
-import path from "path"
-import { BlackData } from "../src/black"
-
-const stage = process.argv[2]
-if (!stage) throw new Error("Stage is required")
-
-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_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_LIMITS ${value} --stage ${stage}`
diff --git a/packages/console/core/script/promote-lite.ts b/packages/console/core/script/promote-limits.ts
index 8fd58c805..f488aba02 100755
--- a/packages/console/core/script/promote-lite.ts
+++ b/packages/console/core/script/promote-limits.ts
@@ -2,7 +2,7 @@
 
 import { $ } from "bun"
 import path from "path"
-import { LiteData } from "../src/lite"
+import { Subscription } from "../src/subscription"
 
 const stage = process.argv[2]
 if (!stage) throw new Error("Stage is required")
@@ -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_LITE_LIMITS"))?.split("=")[1]
-if (!value) throw new Error("ZEN_LITE_LIMITS not found")
+const value = lines.find((line) => line.startsWith("ZEN_LIMITS"))?.split("=")[1]
+if (!value) throw new Error("ZEN_LIMITS not found")
 
 // validate value
-LiteData.validate(JSON.parse(value))
+Subscription.validate(JSON.parse(value))
 
 // update the secret
-await $`bun sst secret set ZEN_LITE_LIMITS ${value} --stage ${stage}`
+await $`bun sst secret set ZEN_LIMITS ${value} --stage ${stage}`
diff --git a/packages/console/core/script/update-black.ts b/packages/console/core/script/update-black.ts
deleted file mode 100755
index 695a5d3ce..000000000
--- a/packages/console/core/script/update-black.ts
+++ /dev/null
@@ -1,28 +0,0 @@
-#!/usr/bin/env bun
-
-import { $ } from "bun"
-import path from "path"
-import os from "os"
-import { BlackData } from "../src/black"
-
-const root = path.resolve(process.cwd(), "..", "..", "..")
-const secrets = await $`bun sst secret list`.cwd(root).text()
-
-// read value
-const lines = secrets.split("\n")
-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`
-const tempFile = Bun.file(path.join(os.tmpdir(), filename))
-await tempFile.write(JSON.stringify(JSON.parse(oldValue), null, 2))
-console.log("tempFile", tempFile.name)
-
-// open temp file in vim and read the file on close
-await $`vim ${tempFile.name}`
-const newValue = JSON.stringify(JSON.parse(await tempFile.text()))
-BlackData.validate(JSON.parse(newValue))
-
-// update the secret
-await $`bun sst secret set ZEN_BLACK_LIMITS ${newValue}`
diff --git a/packages/console/core/script/update-lite.ts b/packages/console/core/script/update-limits.ts
index 2f3e66835..8f2579312 100755
--- a/packages/console/core/script/update-lite.ts
+++ b/packages/console/core/script/update-limits.ts
@@ -3,18 +3,18 @@
 import { $ } from "bun"
 import path from "path"
 import os from "os"
-import { LiteData } from "../src/lite"
+import { Subscription } from "../src/subscription"
 
 const root = path.resolve(process.cwd(), "..", "..", "..")
 const secrets = await $`bun sst secret list`.cwd(root).text()
 
 // read value
 const lines = secrets.split("\n")
-const oldValue = lines.find((line) => line.startsWith("ZEN_LITE_LIMITS"))?.split("=")[1] ?? "{}"
-if (!oldValue) throw new Error("ZEN_LITE_LIMITS not found")
+const oldValue = lines.find((line) => line.startsWith("ZEN_LIMITS"))?.split("=")[1] ?? "{}"
+if (!oldValue) throw new Error("ZEN_LIMITS not found")
 
 // store the prettified json to a temp file
-const filename = `lite-${Date.now()}.json`
+const filename = `limits-${Date.now()}.json`
 const tempFile = Bun.file(path.join(os.tmpdir(), filename))
 await tempFile.write(JSON.stringify(JSON.parse(oldValue), null, 2))
 console.log("tempFile", tempFile.name)
@@ -22,7 +22,7 @@ console.log("tempFile", tempFile.name)
 // open temp file in vim and read the file on close
 await $`vim ${tempFile.name}`
 const newValue = JSON.stringify(JSON.parse(await tempFile.text()))
-LiteData.validate(JSON.parse(newValue))
+Subscription.validate(JSON.parse(newValue))
 
 // update the secret
-await $`bun sst secret set ZEN_LITE_LIMITS ${newValue}`
+await $`bun sst secret set ZEN_LIMITS ${newValue}`