zen: track session in usage

2 files changed, 16 insertions, 10 deletions

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 ce0eb2c6a..c7ca496e6 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
@@ -5,7 +5,7 @@ import { Billing } from "@opencode-ai/console-core/billing.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"
+import { Subscription } from "@opencode-ai/console-core/black.js"
 import { withActor } from "~/context/auth.withActor"
 import { queryBillingInfo } from "../../common"
 import styles from "./black-section.module.css"
@@ -35,12 +35,12 @@ const querySubscription = query(async (workspaceID: string) => {
     return {
       plan: row.subscription.plan,
       useBalance: row.subscription.useBalance ?? false,
-      rollingUsage: Black.analyzeRollingUsage({
+      rollingUsage: Subscription.analyzeRollingUsage({
         plan: row.subscription.plan,
         usage: row.rollingUsage ?? 0,
         timeUpdated: row.timeRollingUpdated ?? new Date(),
       }),
-      weeklyUsage: Black.analyzeWeeklyUsage({
+      weeklyUsage: Subscription.analyzeWeeklyUsage({
         plan: row.subscription.plan,
         usage: row.fixedUsage ?? 0,
         timeUpdated: row.timeFixedUpdated ?? new Date(),
diff --git a/packages/console/app/src/routes/zen/util/handler.ts b/packages/console/app/src/routes/zen/util/handler.ts
index 5f2b51c21..dc10e1bf9 100644
--- a/packages/console/app/src/routes/zen/util/handler.ts
+++ b/packages/console/app/src/routes/zen/util/handler.ts
@@ -9,7 +9,8 @@ import { Billing } from "@opencode-ai/console-core/billing.js"
 import { Actor } from "@opencode-ai/console-core/actor.js"
 import { WorkspaceTable } from "@opencode-ai/console-core/schema/workspace.sql.js"
 import { ZenData } from "@opencode-ai/console-core/model.js"
-import { Black, BlackData } from "@opencode-ai/console-core/black.js"
+import { Subscription } from "@opencode-ai/console-core/subscription.js"
+import { BlackData } from "@opencode-ai/console-core/black.js"
 import { UserTable } from "@opencode-ai/console-core/schema/user.sql.js"
 import { ModelTable } from "@opencode-ai/console-core/schema/model.sql.js"
 import { ProviderTable } from "@opencode-ai/console-core/schema/provider.sql.js"
@@ -196,7 +197,7 @@ export async function handler(
       const costInfo = calculateCost(modelInfo, usageInfo)
       await trialLimiter?.track(usageInfo)
       await rateLimiter?.track()
-      await trackUsage(billingSource, authInfo, modelInfo, providerInfo, usageInfo, costInfo)
+      await trackUsage(sessionId, billingSource, authInfo, modelInfo, providerInfo, usageInfo, costInfo)
       await reload(billingSource, authInfo, costInfo)
 
       const responseConverter = createResponseConverter(providerInfo.format, opts.format)
@@ -246,7 +247,7 @@ export async function handler(
                   const usageInfo = providerInfo.normalizeUsage(usage)
                   const costInfo = calculateCost(modelInfo, usageInfo)
                   await trialLimiter?.track(usageInfo)
-                  await trackUsage(billingSource, authInfo, modelInfo, providerInfo, usageInfo, costInfo)
+                  await trackUsage(sessionId, billingSource, authInfo, modelInfo, providerInfo, usageInfo, costInfo)
                   await reload(billingSource, authInfo, costInfo)
                   cost = calculateOccuredCost(billingSource, costInfo)
                 }
@@ -541,8 +542,9 @@ export async function handler(
 
         // Check weekly limit
         if (sub.fixedUsage && sub.timeFixedUpdated) {
-          const result = Black.analyzeWeeklyUsage({
-            plan,
+          const blackData = BlackData.getLimits({ plan })
+          const result = Subscription.analyzeWeeklyUsage({
+            limit: blackData.fixedLimit,
             usage: sub.fixedUsage,
             timeUpdated: sub.timeFixedUpdated,
           })
@@ -555,8 +557,10 @@ export async function handler(
 
         // Check rolling limit
         if (sub.rollingUsage && sub.timeRollingUpdated) {
-          const result = Black.analyzeRollingUsage({
-            plan,
+          const blackData = BlackData.getLimits({ plan })
+          const result = Subscription.analyzeRollingUsage({
+            limit: blackData.rollingLimit,
+            window: blackData.rollingWindow,
             usage: sub.rollingUsage,
             timeUpdated: sub.timeRollingUpdated,
           })
@@ -687,6 +691,7 @@ export async function handler(
   }
 
   async function trackUsage(
+    sessionId: string,
     billingSource: BillingSource,
     authInfo: AuthInfo,
     modelInfo: ModelInfo,
@@ -734,6 +739,7 @@ export async function handler(
           cacheWrite1hTokens,
           cost,
           keyID: authInfo.apiKeyId,
+          sessionID: sessionId.substring(0, 30),
           enrichment: billingSource === "subscription" ? { plan: "sub" } : undefined,
         }),
         db