Skip to content

Commit e4d35af

Browse files
waleedlatif1waleedlatif1
authored
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
1 parent 1d74ccf commit e4d35af

File tree

4 files changed

+93
-67
lines changed

4 files changed

+93
-67
lines changed

‎apps/sim/lib/auth.ts‎

Lines changed: 18 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -73,6 +73,24 @@ export const auth = betterAuth({
7373
freshAge: 60 * 60, // 1 hour (or set to 0 to disable completely)
7474
},
7575
databaseHooks: {
76+
user: {
77+
create: {
78+
after: async (user) => {
79+
logger.info('[databaseHooks.user.create.after] User created, initializing stats', {
80+
userId: user.id,
81+
})
82+
83+
try {
84+
await handleNewUser(user.id)
85+
} catch (error) {
86+
logger.error('[databaseHooks.user.create.after] Failed to initialize user stats', {
87+
userId: user.id,
88+
error,
89+
})
90+
}
91+
},
92+
},
93+
},
7694
session: {
7795
create: {
7896
before: async (session) => {
@@ -1152,15 +1170,6 @@ export const auth = betterAuth({
11521170
stripeCustomerId: stripeCustomer.id,
11531171
userId: user.id,
11541172
})
1155-
1156-
try {
1157-
await handleNewUser(user.id)
1158-
} catch (error) {
1159-
logger.error('[onCustomerCreate] Failed to handle new user setup', {
1160-
userId: user.id,
1161-
error,
1162-
})
1163-
}
11641173
},
11651174
subscription: {
11661175
enabled: true,

‎apps/sim/lib/billing/calculations/usage-monitor.ts‎

Lines changed: 20 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,6 @@ import { createLogger } from '@/lib/logs/console/logger'
88

99
const logger = createLogger('UsageMonitor')
1010

11-
// Percentage threshold for showing warning
1211
const WARNING_THRESHOLD = 80
1312

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

160-
// Return default values in case of error
159+
// Block execution if we can't determine usage status
160+
logger.error('Cannot determine usage status - blocking execution', {
161+
userId,
162+
error: error instanceof Error ? error.message : String(error),
163+
})
164+
161165
return {
162-
percentUsed: 0,
166+
percentUsed: 100,
163167
isWarning: false,
164-
isExceeded: false,
168+
isExceeded: true, // Block execution when we can't determine status
165169
currentUsage: 0,
166-
limit: 0,
170+
limit: 0, // Zero limit forces blocking
167171
}
168172
}
169173
}
@@ -241,7 +245,6 @@ export async function checkServerSideUsageLimits(userId: string): Promise<{
241245
message?: string
242246
}> {
243247
try {
244-
// If billing is disabled, always allow execution
245248
if (!isBillingEnabled) {
246249
return {
247250
isExceeded: false,
@@ -252,7 +255,6 @@ export async function checkServerSideUsageLimits(userId: string): Promise<{
252255

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

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

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

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

294-
// Be conservative in case of error - allow execution but log the issue
295+
logger.error('Cannot determine usage limits - blocking execution', {
296+
userId,
297+
error: error instanceof Error ? error.message : String(error),
298+
})
299+
295300
return {
296-
isExceeded: false,
301+
isExceeded: true, // Block execution when we can't determine limits
297302
currentUsage: 0,
298-
limit: 0,
299-
message: `Error checking usage limits: ${error instanceof Error ? error.message : String(error)}`,
303+
limit: 0, // Zero limit forces blocking
304+
message:
305+
error instanceof Error && error.message.includes('No user stats record found')
306+
? 'User account not properly initialized. Please contact support.'
307+
: 'Unable to determine usage limits. Execution blocked for security. Please contact support.',
300308
}
301309
}
302310
}

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

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -312,13 +312,15 @@ export async function getUserUsageLimit(userId: string): Promise<number> {
312312
.limit(1)
313313

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

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

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

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

338340
if (orgData[0].orgUsageLimit) {

‎apps/sim/lib/logs/execution/logger.ts‎

Lines changed: 50 additions & 43 deletions
Original file line numberDiff line numberDiff line change
@@ -403,52 +403,59 @@ export class ExecutionLogger implements IExecutionLoggerService {
403403
// Apply cost multiplier only to model costs, not base execution charge
404404
const costToStore = costSummary.baseExecutionCharge + costSummary.modelCost * costMultiplier
405405

406-
// Check if user stats record exists
407-
const userStatsRecords = await db.select().from(userStats).where(eq(userStats.userId, userId))
408-
409-
if (userStatsRecords.length > 0) {
410-
// Update user stats record with trigger-specific increments
411-
const updateFields: any = {
412-
totalTokensUsed: sql`total_tokens_used + ${costSummary.totalTokens}`,
413-
totalCost: sql`total_cost + ${costToStore}`,
414-
currentPeriodCost: sql`current_period_cost + ${costToStore}`, // Track current billing period usage
415-
lastActive: new Date(),
416-
}
417-
418-
// Add trigger-specific increment
419-
switch (trigger) {
420-
case 'manual':
421-
updateFields.totalManualExecutions = sql`total_manual_executions + 1`
422-
break
423-
case 'api':
424-
updateFields.totalApiCalls = sql`total_api_calls + 1`
425-
break
426-
case 'webhook':
427-
updateFields.totalWebhookTriggers = sql`total_webhook_triggers + 1`
428-
break
429-
case 'schedule':
430-
updateFields.totalScheduledExecutions = sql`total_scheduled_executions + 1`
431-
break
432-
case 'chat':
433-
updateFields.totalChatExecutions = sql`total_chat_executions + 1`
434-
break
435-
}
436-
437-
await db.update(userStats).set(updateFields).where(eq(userStats.userId, userId))
406+
// Upsert user stats record - insert if doesn't exist, update if it does
407+
const { getFreeTierLimit } = await import('@/lib/billing/subscriptions/utils')
408+
const defaultLimit = getFreeTierLimit()
409+
410+
const triggerIncrements: any = {}
411+
switch (trigger) {
412+
case 'manual':
413+
triggerIncrements.totalManualExecutions = sql`total_manual_executions + 1`
414+
break
415+
case 'api':
416+
triggerIncrements.totalApiCalls = sql`total_api_calls + 1`
417+
break
418+
case 'webhook':
419+
triggerIncrements.totalWebhookTriggers = sql`total_webhook_triggers + 1`
420+
break
421+
case 'schedule':
422+
triggerIncrements.totalScheduledExecutions = sql`total_scheduled_executions + 1`
423+
break
424+
case 'chat':
425+
triggerIncrements.totalChatExecutions = sql`total_chat_executions + 1`
426+
break
427+
}
438428

439-
logger.debug('Updated user stats record with cost data', {
440-
userId,
441-
trigger,
442-
addedCost: costToStore,
443-
addedTokens: costSummary.totalTokens,
429+
await db
430+
.insert(userStats)
431+
.values({
432+
id: uuidv4(),
433+
userId: userId,
434+
currentUsageLimit: defaultLimit.toString(),
435+
usageLimitUpdatedAt: new Date(),
436+
totalTokensUsed: costSummary.totalTokens,
437+
totalCost: costToStore,
438+
currentPeriodCost: costToStore,
439+
lastActive: new Date(),
440+
...triggerIncrements,
444441
})
445-
} else {
446-
logger.error('User stats record not found - should be created during onboarding', {
447-
userId,
448-
trigger,
442+
.onConflictDoUpdate({
443+
target: userStats.userId,
444+
set: {
445+
totalTokensUsed: sql`total_tokens_used + ${costSummary.totalTokens}`,
446+
totalCost: sql`total_cost + ${costToStore}`,
447+
currentPeriodCost: sql`current_period_cost + ${costToStore}`,
448+
lastActive: new Date(),
449+
...triggerIncrements,
450+
},
449451
})
450-
return // Skip cost tracking if user stats doesn't exist
451-
}
452+
453+
logger.debug('Upserted user stats record with cost data', {
454+
userId,
455+
trigger,
456+
addedCost: costToStore,
457+
addedTokens: costSummary.totalTokens,
458+
})
452459
} catch (error) {
453460
logger.error('Error updating user stats with cost information', {
454461
workflowId,

0 commit comments

Comments
 (0)