diff --git a/apps/sim/lib/auth.ts b/apps/sim/lib/auth.ts index 3e0b1f0064..6d57f54fac 100644 --- 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, diff --git a/apps/sim/lib/billing/calculations/usage-monitor.ts b/apps/sim/lib/billing/calculations/usage-monitor.ts index 4ab0a24ee9..46f0c021c4 100644 --- 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 { 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.', } } } diff --git a/apps/sim/lib/billing/core/usage.ts b/apps/sim/lib/billing/core/usage.ts index 3502f50af2..0b4951c69f 100644 --- 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 { .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 { .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) { diff --git a/apps/sim/lib/logs/execution/logger.ts b/apps/sim/lib/logs/execution/logger.ts index 0fc75752eb..9abbad12f7 100644 --- 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)) - - 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)) + // Upsert user stats record - insert if doesn't exist, update if it does + const { getFreeTierLimit } = await import('@/lib/billing/subscriptions/utils') + const defaultLimit = getFreeTierLimit() + + 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 + } - logger.debug('Updated user stats record with cost data', { - userId, - trigger, - addedCost: costToStore, - addedTokens: costSummary.totalTokens, + 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, }) - } else { - logger.error('User stats record not found - should be created during onboarding', { - userId, - trigger, + .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, + }, }) - return // Skip cost tracking if user stats doesn't exist - } + + 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,