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

Author: waleedlatif1

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

@@ -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/client/utils.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/core/usage.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', {