improvement(usage): bar execution if limits cannot be determined, init user stats record on user creation instead of in stripe plugin (#1399)

* improvement(usage): bar execution if limits cannot be determined, init user stats record on user creation instead of in stripe plugin * upsert user stats record in execution logger
2026-04-06 03:00:16 -04:00 · 2025-09-20 14:08:51 -07:00
parent 1d74ccfeda
commit e4d35afe1f
4 changed files with 94 additions and 68 deletions
--- a/apps/sim/lib/auth.ts
+++ b/apps/sim/lib/auth.ts
@@ -73,6 +73,24 @@ export const auth = betterAuth({
    freshAge: 60 * 60, // 1 hour (or set to 0 to disable completely)
  },
  databaseHooks: {
+    user: {
+      create: {
+        after: async (user) => {
+          logger.info('[databaseHooks.user.create.after] User created, initializing stats', {
+            userId: user.id,
+          })
+
+          try {
+            await handleNewUser(user.id)
+          } catch (error) {
+            logger.error('[databaseHooks.user.create.after] Failed to initialize user stats', {
+              userId: user.id,
+              error,
+            })
+          }
+        },
+      },
+    },
    session: {
      create: {
        before: async (session) => {
@@ -1152,15 +1170,6 @@ export const auth = betterAuth({
                stripeCustomerId: stripeCustomer.id,
                userId: user.id,
              })
-
-              try {
-                await handleNewUser(user.id)
-              } catch (error) {
-                logger.error('[onCustomerCreate] Failed to handle new user setup', {
-                  userId: user.id,
-                  error,
-                })
-              }
            },
            subscription: {
              enabled: true,
--- a/apps/sim/lib/billing/calculations/usage-monitor.ts
+++ b/apps/sim/lib/billing/calculations/usage-monitor.ts
@@ -8,7 +8,6 @@ import { createLogger } from '@/lib/logs/console/logger'

 const logger = createLogger('UsageMonitor')

-// Percentage threshold for showing warning
 const WARNING_THRESHOLD = 80

 interface UsageData {
@@ -157,13 +156,18 @@ export async function checkUsageStatus(userId: string): Promise<UsageData> {
      userId,
    })

-    // Return default values in case of error
+    // Block execution if we can't determine usage status
+    logger.error('Cannot determine usage status - blocking execution', {
+      userId,
+      error: error instanceof Error ? error.message : String(error),
+    })
+
    return {
-      percentUsed: 0,
+      percentUsed: 100,
      isWarning: false,
-      isExceeded: false,
+      isExceeded: true, // Block execution when we can't determine status
      currentUsage: 0,
-      limit: 0,
+      limit: 0, // Zero limit forces blocking
    }
  }
 }
@@ -241,7 +245,6 @@ export async function checkServerSideUsageLimits(userId: string): Promise<{
  message?: string
 }> {
  try {
-    // If billing is disabled, always allow execution
    if (!isBillingEnabled) {
      return {
        isExceeded: false,
@@ -252,7 +255,6 @@ export async function checkServerSideUsageLimits(userId: string): Promise<{

    logger.info('Server-side checking usage limits for user', { userId })

-    // Hard block if billing is flagged as blocked
    const stats = await db
      .select({
        blocked: userStats.billingBlocked,
@@ -274,7 +276,6 @@ export async function checkServerSideUsageLimits(userId: string): Promise<{
      }
    }

-    // Get usage data using the same function we use for client-side
    const usageData = await checkUsageStatus(userId)

    return {
@@ -291,12 +292,19 @@ export async function checkServerSideUsageLimits(userId: string): Promise<{
      userId,
    })

-    // Be conservative in case of error - allow execution but log the issue
+    logger.error('Cannot determine usage limits - blocking execution', {
+      userId,
+      error: error instanceof Error ? error.message : String(error),
+    })
+
    return {
-      isExceeded: false,
+      isExceeded: true, // Block execution when we can't determine limits
      currentUsage: 0,
-      limit: 0,
-      message: `Error checking usage limits: ${error instanceof Error ? error.message : String(error)}`,
+      limit: 0, // Zero limit forces blocking
+      message:
+        error instanceof Error && error.message.includes('No user stats record found')
+          ? 'User account not properly initialized. Please contact support.'
+          : 'Unable to determine usage limits. Execution blocked for security. Please contact support.',
    }
  }
 }
--- a/apps/sim/lib/billing/core/usage.ts
+++ b/apps/sim/lib/billing/core/usage.ts
@@ -312,13 +312,15 @@ export async function getUserUsageLimit(userId: string): Promise<number> {
      .limit(1)

    if (userStatsQuery.length === 0) {
-      throw new Error(`User stats not found for userId: ${userId}`)
+      throw new Error(
+        `No user stats record found for userId: ${userId}. User must be properly initialized before execution.`
+      )
    }

    // Individual limits should never be null for free/pro users
    if (!userStatsQuery[0].currentUsageLimit) {
      throw new Error(
-        `Invalid null usage limit for ${subscription?.plan || 'free'} user: ${userId}`
+        `Invalid null usage limit for ${subscription?.plan || 'free'} user: ${userId}. User stats must be properly initialized.`
      )
    }

@@ -332,7 +334,7 @@ export async function getUserUsageLimit(userId: string): Promise<number> {
    .limit(1)

  if (orgData.length === 0) {
-    throw new Error(`Organization not found: ${subscription.referenceId}`)
+    throw new Error(`Organization not found: ${subscription.referenceId} for user: ${userId}`)
  }

  if (orgData[0].orgUsageLimit) {
--- a/apps/sim/lib/logs/execution/logger.ts
+++ b/apps/sim/lib/logs/execution/logger.ts
@@ -403,52 +403,59 @@ export class ExecutionLogger implements IExecutionLoggerService {
      // Apply cost multiplier only to model costs, not base execution charge
      const costToStore = costSummary.baseExecutionCharge + costSummary.modelCost * costMultiplier

-      // Check if user stats record exists
-      const userStatsRecords = await db.select().from(userStats).where(eq(userStats.userId, userId))
+      // Upsert user stats record - insert if doesn't exist, update if it does
+      const { getFreeTierLimit } = await import('@/lib/billing/subscriptions/utils')
+      const defaultLimit = getFreeTierLimit()

-      if (userStatsRecords.length > 0) {
-        // Update user stats record with trigger-specific increments
-        const updateFields: any = {
-          totalTokensUsed: sql`total_tokens_used + ${costSummary.totalTokens}`,
-          totalCost: sql`total_cost + ${costToStore}`,
-          currentPeriodCost: sql`current_period_cost + ${costToStore}`, // Track current billing period usage
-          lastActive: new Date(),
-        }
-
-        // Add trigger-specific increment
-        switch (trigger) {
-          case 'manual':
-            updateFields.totalManualExecutions = sql`total_manual_executions + 1`
-            break
-          case 'api':
-            updateFields.totalApiCalls = sql`total_api_calls + 1`
-            break
-          case 'webhook':
-            updateFields.totalWebhookTriggers = sql`total_webhook_triggers + 1`
-            break
-          case 'schedule':
-            updateFields.totalScheduledExecutions = sql`total_scheduled_executions + 1`
-            break
-          case 'chat':
-            updateFields.totalChatExecutions = sql`total_chat_executions + 1`
-            break
-        }
-
-        await db.update(userStats).set(updateFields).where(eq(userStats.userId, userId))
-
-        logger.debug('Updated user stats record with cost data', {
-          userId,
-          trigger,
-          addedCost: costToStore,
-          addedTokens: costSummary.totalTokens,
-        })
-      } else {
-        logger.error('User stats record not found - should be created during onboarding', {
-          userId,
-          trigger,
-        })
-        return // Skip cost tracking if user stats doesn't exist
+      const triggerIncrements: any = {}
+      switch (trigger) {
+        case 'manual':
+          triggerIncrements.totalManualExecutions = sql`total_manual_executions + 1`
+          break
+        case 'api':
+          triggerIncrements.totalApiCalls = sql`total_api_calls + 1`
+          break
+        case 'webhook':
+          triggerIncrements.totalWebhookTriggers = sql`total_webhook_triggers + 1`
+          break
+        case 'schedule':
+          triggerIncrements.totalScheduledExecutions = sql`total_scheduled_executions + 1`
+          break
+        case 'chat':
+          triggerIncrements.totalChatExecutions = sql`total_chat_executions + 1`
+          break
      }
+
+      await db
+        .insert(userStats)
+        .values({
+          id: uuidv4(),
+          userId: userId,
+          currentUsageLimit: defaultLimit.toString(),
+          usageLimitUpdatedAt: new Date(),
+          totalTokensUsed: costSummary.totalTokens,
+          totalCost: costToStore,
+          currentPeriodCost: costToStore,
+          lastActive: new Date(),
+          ...triggerIncrements,
+        })
+        .onConflictDoUpdate({
+          target: userStats.userId,
+          set: {
+            totalTokensUsed: sql`total_tokens_used + ${costSummary.totalTokens}`,
+            totalCost: sql`total_cost + ${costToStore}`,
+            currentPeriodCost: sql`current_period_cost + ${costToStore}`,
+            lastActive: new Date(),
+            ...triggerIncrements,
+          },
+        })
+
+      logger.debug('Upserted user stats record with cost data', {
+        userId,
+        trigger,
+        addedCost: costToStore,
+        addedTokens: costSummary.totalTokens,
+      })
    } catch (error) {
      logger.error('Error updating user stats with cost information', {
        workflowId,