improvement(billing): move overage calculations out of txes (#4595)

* improvement(billing): move calc subscription overage out of tx

* fix double billing risk

* address comments

* address comments

* share timeout const

Author: icecrasher321

apps/sim/lib/billing/constants.ts

@@ -34,6 +34,11 @@ export const SEARCH_TOOL_COST = 0.01
  */
 export const DEFAULT_OVERAGE_THRESHOLD = 100
 
+/**
+ * Maximum time to wait on billing coordination row locks before retrying later.
+ */
+export const BILLING_LOCK_TIMEOUT_MS = 5_000
+
 /**
  * Available credit tiers. Each tier maps a credit amount to the underlying dollar cost.
  * 1 credit = $0.005, so credits = dollars * 200.

apps/sim/lib/billing/organizations/membership.ts

@@ -926,34 +926,6 @@ export async function removeUserFromOrganization(
         )
       }
 
-      let capturedUsage = 0
-      if (!skipBillingLogic) {
-        const [departingUserStats] = await tx
-          .select({ currentPeriodCost: userStats.currentPeriodCost })
-          .from(userStats)
-          .where(eq(userStats.userId, userId))
-          .limit(1)
-
-        if (departingUserStats?.currentPeriodCost) {
-          const usage = toNumber(toDecimal(departingUserStats.currentPeriodCost))
-          if (usage > 0) {
-            await tx
-              .update(organization)
-              .set({
-                departedMemberUsage: sql`${organization.departedMemberUsage} + ${usage}`,
-              })
-              .where(eq(organization.id, organizationId))
-
-            await tx
-              .update(userStats)
-              .set({ currentPeriodCost: '0' })
-              .where(eq(userStats.userId, userId))
-
-            capturedUsage = usage
-          }
-        }
-      }
-
       const [targetUser] = await tx
         .select({ email: user.email })
         .from(user)
@@ -979,7 +951,44 @@ export async function removeUserFromOrganization(
         .from(workspace)
         .where(eq(workspace.organizationId, organizationId))
 
+      const captureDepartedUsage = async () => {
+        if (skipBillingLogic) return 0
+
+        await tx
+          .select({ id: organization.id })
+          .from(organization)
+          .where(eq(organization.id, organizationId))
+          .for('update')
+          .limit(1)
+
+        const [departingUserStats] = await tx
+          .select({ currentPeriodCost: userStats.currentPeriodCost })
+          .from(userStats)
+          .where(eq(userStats.userId, userId))
+          .for('update')
+          .limit(1)
+
+        const usage = toNumber(toDecimal(departingUserStats?.currentPeriodCost))
+        if (usage <= 0) return 0
+
+        await tx
+          .update(organization)
+          .set({
+            departedMemberUsage: sql`${organization.departedMemberUsage} + ${usage}`,
+          })
+          .where(eq(organization.id, organizationId))
+
+        await tx
+          .update(userStats)
+          .set({ currentPeriodCost: '0' })
+          .where(eq(userStats.userId, userId))
+
+        return usage
+      }
+
       if (orgWorkspaces.length === 0) {
+        const capturedUsage = await captureDepartedUsage()
+
         return {
           workspaceIdsToRevoke: [] as string[],
           usageCaptured: capturedUsage,
@@ -1022,6 +1031,7 @@ export async function removeUserFromOrganization(
         workspaceIds,
         userId,
       })
+      const capturedUsage = await captureDepartedUsage()
 
       return {
         workspaceIdsToRevoke: deletedPerms.map((row) => row.entityId),