fix(billing): handle missing userStats and prevent crashes (#2956)

* fix(billing): handle missing userStats and prevent crashes

* fix(billing): correct import path for getFilledPillColor

* fix(billing): add Number.isFinite check to lastPeriodCost

apps/sim/app/workspace/[workspaceId]/w/components/sidebar/components/settings-modal/components/subscription/subscription.tsx

@@ -6,8 +6,8 @@ import { useParams } from 'next/navigation'
 import { Combobox, Label, Switch, Tooltip } from '@/components/emcn'
 import { Skeleton } from '@/components/ui'
 import { useSession } from '@/lib/auth/auth-client'
+import { USAGE_THRESHOLDS } from '@/lib/billing/client/consts'
 import { useSubscriptionUpgrade } from '@/lib/billing/client/upgrade'
-import { USAGE_THRESHOLDS } from '@/lib/billing/client/usage-visualization'
 import { getEffectiveSeats } from '@/lib/billing/subscriptions/utils'
 import { cn } from '@/lib/core/utils/cn'
 import { getBaseUrl } from '@/lib/core/utils/urls'

apps/sim/app/workspace/[workspaceId]/w/components/sidebar/components/settings-modal/components/usage-header/usage-header.tsx

@@ -2,11 +2,7 @@
 
 import type { ReactNode } from 'react'
 import { Badge } from '@/components/emcn'
-import {
-  getFilledPillColor,
-  USAGE_PILL_COLORS,
-  USAGE_THRESHOLDS,
-} from '@/lib/billing/client/usage-visualization'
+import { getFilledPillColor, USAGE_PILL_COLORS, USAGE_THRESHOLDS } from '@/lib/billing/client'
 
 const PILL_COUNT = 5
 

apps/sim/app/workspace/[workspaceId]/w/components/sidebar/components/usage-indicator/usage-indicator.tsx

@@ -5,13 +5,14 @@ import { createLogger } from '@sim/logger'
 import { useQueryClient } from '@tanstack/react-query'
 import { Badge } from '@/components/emcn'
 import { Skeleton } from '@/components/ui'
+import { USAGE_PILL_COLORS, USAGE_THRESHOLDS } from '@/lib/billing/client/consts'
 import { useSubscriptionUpgrade } from '@/lib/billing/client/upgrade'
 import {
+  getBillingStatus,
   getFilledPillColor,
-  USAGE_PILL_COLORS,
-  USAGE_THRESHOLDS,
-} from '@/lib/billing/client/usage-visualization'
-import { getBillingStatus, getSubscriptionStatus, getUsage } from '@/lib/billing/client/utils'
+  getSubscriptionStatus,
+  getUsage,
+} from '@/lib/billing/client/utils'
 import { useContextMenu } from '@/app/workspace/[workspaceId]/w/components/sidebar/hooks'
 import { useSocket } from '@/app/workspace/providers/socket-provider'
 import { subscriptionKeys, useSubscriptionData } from '@/hooks/queries/subscription'

apps/sim/lib/auth/auth.ts

@@ -256,6 +256,17 @@ export const auth = betterAuth({
           return { data: account }
         },
         after: async (account) => {
+          try {
+            const { ensureUserStatsExists } = await import('@/lib/billing/core/usage')
+            await ensureUserStatsExists(account.userId)
+          } catch (error) {
+            logger.error('[databaseHooks.account.create.after] Failed to ensure user stats', {
+              userId: account.userId,
+              accountId: account.id,
+              error,
+            })
+          }
+
           if (account.providerId === 'salesforce') {
             const updates: {
               accessTokenExpiresAt?: Date
@@ -462,7 +473,6 @@ export const auth = betterAuth({
   },
   emailVerification: {
     autoSignInAfterVerification: true,
-    // onEmailVerification is called by the emailOTP plugin when email is verified via OTP
     onEmailVerification: async (user) => {
       if (isHosted && user.email) {
         try {

apps/sim/lib/billing/client/consts.ts

@@ -0,0 +1,28 @@
+/**
+ * Number of pills to display in usage indicators.
+ */
+export const USAGE_PILL_COUNT = 8
+
+/**
+ * Usage percentage thresholds for visual states.
+ */
+export const USAGE_THRESHOLDS = {
+  /** Warning threshold (yellow/orange state) */
+  WARNING: 75,
+  /** Critical threshold (red state) */
+  CRITICAL: 90,
+} as const
+
+/**
+ * Color values for usage pill states using CSS variables
+ */
+export const USAGE_PILL_COLORS = {
+  /** Unfilled pill color (gray) */
+  UNFILLED: 'var(--surface-7)',
+  /** Normal filled pill color (blue) */
+  FILLED: 'var(--brand-secondary)',
+  /** Warning state pill color (yellow/orange) */
+  WARNING: 'var(--warning)',
+  /** Critical/limit reached pill color (red) */
+  AT_LIMIT: 'var(--text-error)',
+} as const

apps/sim/lib/billing/client/index.ts

@@ -1,3 +1,7 @@
+export {
+  USAGE_PILL_COLORS,
+  USAGE_THRESHOLDS,
+} from './consts'
 export type {
   BillingStatus,
   SubscriptionData,
@@ -8,6 +12,7 @@ export {
   canUpgrade,
   getBillingStatus,
   getDaysRemainingInPeriod,
+  getFilledPillColor,
   getRemainingBudget,
   getSubscriptionStatus,
   getUsage,

apps/sim/lib/billing/client/usage-visualization.ts

@@ -1,150 +0,0 @@
-/**
- * Shared utilities for consistent usage visualization across the application.
- *
- * This module provides a single source of truth for how usage metrics are
- * displayed visually through "pills" or progress indicators.
- */
-
-/**
- * Number of pills to display in usage indicators.
- *
- * Using 8 pills provides:
- * - 12.5% granularity per pill
- * - Good balance between precision and visual clarity
- * - Consistent representation across panel and settings
- */
-export const USAGE_PILL_COUNT = 8
-
-/**
- * Usage percentage thresholds for visual states.
- */
-export const USAGE_THRESHOLDS = {
-  /** Warning threshold (yellow/orange state) */
-  WARNING: 75,
-  /** Critical threshold (red state) */
-  CRITICAL: 90,
-} as const
-
-/**
- * Color values for usage pill states using CSS variables
- */
-export const USAGE_PILL_COLORS = {
-  /** Unfilled pill color (gray) */
-  UNFILLED: 'var(--surface-7)',
-  /** Normal filled pill color (blue) */
-  FILLED: 'var(--brand-secondary)',
-  /** Warning state pill color (yellow/orange) */
-  WARNING: 'var(--warning)',
-  /** Critical/limit reached pill color (red) */
-  AT_LIMIT: 'var(--text-error)',
-} as const
-
-/**
- * Calculate the number of filled pills based on usage percentage.
- *
- * Uses Math.ceil() to ensure even minimal usage (0.01%) shows visual feedback.
- * This provides better UX by making it clear that there is some usage, even if small.
- *
- * @param percentUsed - The usage percentage (0-100). Can be a decimal (e.g., 0.315 for 0.315%)
- * @returns Number of pills that should be filled (0 to USAGE_PILL_COUNT)
- *
- * @example
- * calculateFilledPills(0.315)  // Returns 1 (shows feedback for 0.315% usage)
- * calculateFilledPills(50)     // Returns 4 (50% of 8 pills)
- * calculateFilledPills(100)    // Returns 8 (completely filled)
- * calculateFilledPills(150)    // Returns 8 (clamped to maximum)
- */
-export function calculateFilledPills(percentUsed: number): number {
-  // Clamp percentage to valid range [0, 100]
-  const safePercent = Math.min(Math.max(percentUsed, 0), 100)
-
-  // Calculate filled pills using ceil to show feedback for any usage
-  return Math.ceil((safePercent / 100) * USAGE_PILL_COUNT)
-}
-
-/**
- * Determine if usage has reached the limit (all pills filled).
- *
- * @param percentUsed - The usage percentage (0-100)
- * @returns true if all pills should be filled (at or over limit)
- */
-export function isUsageAtLimit(percentUsed: number): boolean {
-  return calculateFilledPills(percentUsed) >= USAGE_PILL_COUNT
-}
-
-/**
- * Get the appropriate color for a pill based on its state.
- *
- * @param isFilled - Whether this pill should be filled
- * @param isAtLimit - Whether usage has reached the limit
- * @returns CSS color value
- */
-export function getPillColor(isFilled: boolean, isAtLimit: boolean): string {
-  if (!isFilled) return USAGE_PILL_COLORS.UNFILLED
-  if (isAtLimit) return USAGE_PILL_COLORS.AT_LIMIT
-  return USAGE_PILL_COLORS.FILLED
-}
-
-/**
- * Get the appropriate filled pill color based on usage thresholds.
- *
- * @param isCritical - Whether usage is at critical level (blocked or >= 90%)
- * @param isWarning - Whether usage is at warning level (>= 75% but < critical)
- * @returns CSS color value for filled pills
- */
-export function getFilledPillColor(isCritical: boolean, isWarning: boolean): string {
-  if (isCritical) return USAGE_PILL_COLORS.AT_LIMIT
-  if (isWarning) return USAGE_PILL_COLORS.WARNING
-  return USAGE_PILL_COLORS.FILLED
-}
-
-/**
- * Determine usage state based on percentage and blocked status.
- *
- * @param percentUsed - The usage percentage (0-100)
- * @param isBlocked - Whether the account is blocked
- * @returns Object containing isCritical and isWarning flags
- */
-export function getUsageState(
-  percentUsed: number,
-  isBlocked = false
-): { isCritical: boolean; isWarning: boolean } {
-  const isCritical = isBlocked || percentUsed >= USAGE_THRESHOLDS.CRITICAL
-  const isWarning = !isCritical && percentUsed >= USAGE_THRESHOLDS.WARNING
-  return { isCritical, isWarning }
-}
-
-/**
- * Generate an array of pill states for rendering.
- *
- * @param percentUsed - The usage percentage (0-100)
- * @param isBlocked - Whether the account is blocked
- * @returns Array of pill states with colors
- *
- * @example
- * const pills = generatePillStates(50)
- * pills.forEach((pill, index) => (
- *   <Pill key={index} color={pill.color} filled={pill.filled} />
- * ))
- */
-export function generatePillStates(
-  percentUsed: number,
-  isBlocked = false
-): Array<{
-  filled: boolean
-  color: string
-  index: number
-}> {
-  const filledCount = calculateFilledPills(percentUsed)
-  const { isCritical, isWarning } = getUsageState(percentUsed, isBlocked)
-  const filledColor = getFilledPillColor(isCritical, isWarning)
-
-  return Array.from({ length: USAGE_PILL_COUNT }, (_, index) => {
-    const filled = index < filledCount
-    return {
-      filled,
-      color: filled ? filledColor : USAGE_PILL_COLORS.UNFILLED,
-      index,
-    }
-  })
-}

apps/sim/lib/billing/client/utils.ts

@@ -4,6 +4,7 @@
  */
 
 import { DEFAULT_FREE_CREDITS } from '@/lib/billing/constants'
+import { USAGE_PILL_COLORS } from './consts'
 import type { BillingStatus, SubscriptionData, UsageData } from './types'
 
 const defaultUsage: UsageData = {
@@ -36,9 +37,35 @@ export function getSubscriptionStatus(subscriptionData: SubscriptionData | null
 
 /**
  * Get usage data from subscription data
+ * Validates and sanitizes all numeric values to prevent crashes from malformed data
  */
 export function getUsage(subscriptionData: SubscriptionData | null | undefined): UsageData {
-  return subscriptionData?.usage ?? defaultUsage
+  const usage = subscriptionData?.usage
+
+  if (!usage) {
+    return defaultUsage
+  }
+
+  return {
+    current:
+      typeof usage.current === 'number' && Number.isFinite(usage.current) ? usage.current : 0,
+    limit:
+      typeof usage.limit === 'number' && Number.isFinite(usage.limit)
+        ? usage.limit
+        : DEFAULT_FREE_CREDITS,
+    percentUsed:
+      typeof usage.percentUsed === 'number' && Number.isFinite(usage.percentUsed)
+        ? usage.percentUsed
+        : 0,
+    isWarning: Boolean(usage.isWarning),
+    isExceeded: Boolean(usage.isExceeded),
+    billingPeriodStart: usage.billingPeriodStart ?? null,
+    billingPeriodEnd: usage.billingPeriodEnd ?? null,
+    lastPeriodCost:
+      typeof usage.lastPeriodCost === 'number' && Number.isFinite(usage.lastPeriodCost)
+        ? usage.lastPeriodCost
+        : 0,
+  }
 }
 
 /**
@@ -100,3 +127,16 @@ export function canUpgrade(subscriptionData: SubscriptionData | null | undefined
   const status = getSubscriptionStatus(subscriptionData)
   return status.plan === 'free' || status.plan === 'pro'
 }
+
+/**
+ * Get the appropriate filled pill color based on usage thresholds.
+ *
+ * @param isCritical - Whether usage is at critical level (blocked or >= 90%)
+ * @param isWarning - Whether usage is at warning level (>= 75% but < critical)
+ * @returns CSS color value for filled pills
+ */
+export function getFilledPillColor(isCritical: boolean, isWarning: boolean): string {
+  if (isCritical) return USAGE_PILL_COLORS.AT_LIMIT
+  if (isWarning) return USAGE_PILL_COLORS.WARNING
+  return USAGE_PILL_COLORS.FILLED
+}

apps/sim/lib/billing/core/usage.ts

@@ -96,11 +96,32 @@ export async function handleNewUser(userId: string): Promise<void> {
   }
 }
 
+/**
+ * Ensures a userStats record exists for a user.
+ * Creates one with default values if missing.
+ * This is a fallback for cases where the user.create.after hook didn't fire
+ * (e.g., OAuth account linking to existing users).
+ *
+ */
+export async function ensureUserStatsExists(userId: string): Promise<void> {
+  await db
+    .insert(userStats)
+    .values({
+      id: crypto.randomUUID(),
+      userId: userId,
+      currentUsageLimit: getFreeTierLimit().toString(),
+      usageLimitUpdatedAt: new Date(),
+    })
+    .onConflictDoNothing({ target: userStats.userId })
+}
+
 /**
  * Get comprehensive usage data for a user
  */
 export async function getUserUsageData(userId: string): Promise<UsageData> {
   try {
+    await ensureUserStatsExists(userId)
+
     const [userStatsData, subscription] = await Promise.all([
       db.select().from(userStats).where(eq(userStats.userId, userId)).limit(1),
       getHighestPrioritySubscription(userId),

apps/sim/lib/billing/organizations/membership.ts

@@ -16,6 +16,7 @@ import {
 } from '@sim/db/schema'
 import { createLogger } from '@sim/logger'
 import { and, eq, sql } from 'drizzle-orm'
+import { syncUsageLimitsFromSubscription } from '@/lib/billing/core/usage'
 import { requireStripeClient } from '@/lib/billing/stripe-client'
 import { validateSeatAvailability } from '@/lib/billing/validation/seat-management'
 
@@ -556,6 +557,8 @@ export async function removeUserFromOrganization(
           const restoreResult = await restoreUserProSubscription(userId)
           billingActions.proRestored = restoreResult.restored
           billingActions.usageRestored = restoreResult.usageRestored
+
+          await syncUsageLimitsFromSubscription(userId)
         }
       } catch (postRemoveError) {
         logger.error('Post-removal personal Pro restore check failed', {