address comments

Author: icecrasher321

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),

apps/sim/lib/billing/threshold-billing.ts

@@ -601,7 +601,7 @@ function buildOrganizationUsageSnapshot(
       )
       .join('|'),
     pooledCurrentPeriodCost,
-    departedMemberUsage: toNumber(toDecimal(sortedRows[0]?.departedMemberUsage)),
+    departedMemberUsage: toNumber(toDecimal(owner.departedMemberUsage)),
   }
 }
 
@@ -612,8 +612,6 @@ function organizationUsageSnapshotMatches(
   return (
     expected.ownerId === actual.ownerId &&
     expected.memberSignature === actual.memberSignature &&
-    Math.abs(expected.pooledCurrentPeriodCost - actual.pooledCurrentPeriodCost) <=
-      USAGE_TOTAL_EPSILON &&
     Math.abs(expected.departedMemberUsage - actual.departedMemberUsage) <= USAGE_TOTAL_EPSILON
   )
 }