Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -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'
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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'
Expand Down
12 changes: 11 additions & 1 deletion apps/sim/lib/auth/auth.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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 {
Expand Down
28 changes: 28 additions & 0 deletions apps/sim/lib/billing/client/consts.ts
Original file line number Diff line number Diff line change
@@ -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
5 changes: 5 additions & 0 deletions apps/sim/lib/billing/client/index.ts
Original file line number Diff line number Diff line change
@@ -1,3 +1,7 @@
export {
USAGE_PILL_COLORS,
USAGE_THRESHOLDS,
} from './consts'
export type {
BillingStatus,
SubscriptionData,
Expand All @@ -8,6 +12,7 @@ export {
canUpgrade,
getBillingStatus,
getDaysRemainingInPeriod,
getFilledPillColor,
getRemainingBudget,
getSubscriptionStatus,
getUsage,
Expand Down
150 changes: 0 additions & 150 deletions apps/sim/lib/billing/client/usage-visualization.ts

This file was deleted.

42 changes: 41 additions & 1 deletion apps/sim/lib/billing/client/utils.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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 = {
Expand Down Expand Up @@ -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,
}
}

/**
Expand Down Expand Up @@ -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
}
21 changes: 21 additions & 0 deletions apps/sim/lib/billing/core/usage.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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),
Expand Down
3 changes: 3 additions & 0 deletions apps/sim/lib/billing/organizations/membership.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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'

Expand Down Expand Up @@ -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', {
Expand Down