fix(subscription): incomplete team subscription race condition (#2381)

icecrasher321 · web-flow · commit e43afc8b6ca7 · 2025-12-15T11:36:08.000-08:00
diff --git a/apps/sim/app/api/organizations/route.ts b/apps/sim/app/api/organizations/route.ts
@@ -16,7 +16,6 @@ export async function GET() {
       return NextResponse.json({ error: 'Unauthorized' }, { status: 401 })
     }
 
-    // Get organizations where user is owner or admin
     const userOrganizations = await db
       .select({
         id: organization.id,
@@ -32,8 +31,15 @@ export async function GET() {
         )
       )
 
+    const anyMembership = await db
+      .select({ id: member.id })
+      .from(member)
+      .where(eq(member.userId, session.user.id))
+      .limit(1)
+
     return NextResponse.json({
       organizations: userOrganizations,
+      isMemberOfAnyOrg: anyMembership.length > 0,
     })
   } catch (error) {
     logger.error('Failed to fetch organizations', {
diff --git a/apps/sim/lib/auth/auth.ts b/apps/sim/lib/auth/auth.ts
@@ -24,7 +24,10 @@ import {
 import { sendPlanWelcomeEmail } from '@/lib/billing'
 import { authorizeSubscriptionReference } from '@/lib/billing/authorization'
 import { handleNewUser } from '@/lib/billing/core/usage'
-import { syncSubscriptionUsageLimits } from '@/lib/billing/organization'
+import {
+  ensureOrganizationForTeamSubscription,
+  syncSubscriptionUsageLimits,
+} from '@/lib/billing/organization'
 import { getPlans } from '@/lib/billing/plans'
 import { syncSeatsFromStripeQuantity } from '@/lib/billing/validation/seat-management'
 import { handleChargeDispute, handleDisputeClosed } from '@/lib/billing/webhooks/disputes'
@@ -2021,11 +2024,14 @@ export const auth = betterAuth({
                   status: subscription.status,
                 })
 
-                await handleSubscriptionCreated(subscription)
+                const resolvedSubscription =
+                  await ensureOrganizationForTeamSubscription(subscription)
+
+                await handleSubscriptionCreated(resolvedSubscription)
 
-                await syncSubscriptionUsageLimits(subscription)
+                await syncSubscriptionUsageLimits(resolvedSubscription)
 
-                await sendPlanWelcomeEmail(subscription)
+                await sendPlanWelcomeEmail(resolvedSubscription)
               },
               onSubscriptionUpdate: async ({
                 event,
@@ -2040,40 +2046,42 @@ export const auth = betterAuth({
                   plan: subscription.plan,
                 })
 
+                const resolvedSubscription =
+                  await ensureOrganizationForTeamSubscription(subscription)
+
                 try {
-                  await syncSubscriptionUsageLimits(subscription)
+                  await syncSubscriptionUsageLimits(resolvedSubscription)
                 } catch (error) {
                   logger.error('[onSubscriptionUpdate] Failed to sync usage limits', {
-                    subscriptionId: subscription.id,
-                    referenceId: subscription.referenceId,
+                    subscriptionId: resolvedSubscription.id,
+                    referenceId: resolvedSubscription.referenceId,
                     error,
                   })
                 }
 
-                // Sync seat count from Stripe subscription quantity for team plans
-                if (subscription.plan === 'team') {
+                if (resolvedSubscription.plan === 'team') {
                   try {
                     const stripeSubscription = event.data.object as Stripe.Subscription
                     const quantity = stripeSubscription.items?.data?.[0]?.quantity || 1
 
                     const result = await syncSeatsFromStripeQuantity(
-                      subscription.id,
-                      subscription.seats,
+                      resolvedSubscription.id,
+                      resolvedSubscription.seats ?? null,
                       quantity
                     )
 
                     if (result.synced) {
                       logger.info('[onSubscriptionUpdate] Synced seat count from Stripe', {
-                        subscriptionId: subscription.id,
-                        referenceId: subscription.referenceId,
+                        subscriptionId: resolvedSubscription.id,
+                        referenceId: resolvedSubscription.referenceId,
                         previousSeats: result.previousSeats,
                         newSeats: result.newSeats,
                       })
                     }
                   } catch (error) {
                     logger.error('[onSubscriptionUpdate] Failed to sync seat count', {
-                      subscriptionId: subscription.id,
-                      referenceId: subscription.referenceId,
+                      subscriptionId: resolvedSubscription.id,
+                      referenceId: resolvedSubscription.referenceId,
                       error,
                     })
                   }
diff --git a/apps/sim/lib/billing/client/upgrade.ts b/apps/sim/lib/billing/client/upgrade.ts
@@ -12,9 +12,6 @@ const CONSTANTS = {
   INITIAL_TEAM_SEATS: 1,
 } as const
 
-/**
- * Handles organization creation for team plans and proper referenceId management
- */
 export function useSubscriptionUpgrade() {
   const { data: session } = useSession()
   const betterAuthSubscription = useSubscription()
@@ -40,83 +37,43 @@ export function useSubscriptionUpgrade() {
 
       let referenceId = userId
 
-      // For team plans, create organization first and use its ID as referenceId
       if (targetPlan === 'team') {
         try {
-          // Check if user already has an organization where they are owner/admin
           const orgsResponse = await fetch('/api/organizations')
-          if (orgsResponse.ok) {
-            const orgsData = await orgsResponse.json()
-            const existingOrg = orgsData.organizations?.find(
-              (org: any) => org.role === 'owner' || org.role === 'admin'
-            )
-
-            if (existingOrg) {
-              logger.info('Using existing organization for team plan upgrade', {
-                userId,
-                organizationId: existingOrg.id,
-              })
-              referenceId = existingOrg.id
-            }
+          if (!orgsResponse.ok) {
+            throw new Error('Failed to check organization status')
           }
 
-          // Only create new organization if no suitable one exists
-          if (referenceId === userId) {
-            logger.info('Creating organization for team plan upgrade', {
-              userId,
-            })
-
-            const response = await fetch('/api/organizations', {
-              method: 'POST',
-              headers: {
-                'Content-Type': 'application/json',
-              },
-            })
-
-            if (!response.ok) {
-              const errorData = await response.json().catch(() => ({}))
-              if (response.status === 409) {
-                throw new Error(
-                  'You are already a member of an organization. Please leave it or ask an admin to upgrade.'
-                )
-              }
-              throw new Error(
-                errorData.message || `Failed to create organization: ${response.statusText}`
-              )
-            }
-            const result = await response.json()
+          const orgsData = await orgsResponse.json()
+          const existingOrg = orgsData.organizations?.find(
+            (org: any) => org.role === 'owner' || org.role === 'admin'
+          )
 
-            logger.info('Organization API response', {
-              result,
-              success: result.success,
-              organizationId: result.organizationId,
+          if (existingOrg) {
+            logger.info('Using existing organization for team plan upgrade', {
+              userId,
+              organizationId: existingOrg.id,
             })
+            referenceId = existingOrg.id
 
-            if (!result.success || !result.organizationId) {
-              throw new Error('Failed to create organization for team plan')
+            try {
+              await client.organization.setActive({ organizationId: referenceId })
+              logger.info('Set organization as active', { organizationId: referenceId })
+            } catch (error) {
+              logger.warn('Failed to set organization as active, proceeding with upgrade', {
+                organizationId: referenceId,
+                error: error instanceof Error ? error.message : 'Unknown error',
+              })
             }
-
-            referenceId = result.organizationId
-          }
-
-          // Set the organization as active so Better Auth recognizes it
-          try {
-            await client.organization.setActive({ organizationId: referenceId })
-
-            logger.info('Set organization as active', {
-              organizationId: referenceId,
-              oldReferenceId: userId,
-              newReferenceId: referenceId,
-            })
-          } catch (error) {
-            logger.warn('Failed to set organization as active, but proceeding with upgrade', {
-              organizationId: referenceId,
-              error: error instanceof Error ? error.message : 'Unknown error',
-            })
-            // Continue with upgrade even if setting active fails
+          } else if (orgsData.isMemberOfAnyOrg) {
+            throw new Error(
+              'You are already a member of an organization. Please leave it or ask an admin to upgrade.'
+            )
+          } else {
+            logger.info('Will create organization after payment succeeds', { userId })
           }
         } catch (error) {
-          logger.error('Failed to prepare organization for team plan', error)
+          logger.error('Failed to prepare for team plan upgrade', error)
           throw error instanceof Error
             ? error
             : new Error('Failed to prepare team workspace. Please try again or contact support.')
@@ -134,23 +91,17 @@ export function useSubscriptionUpgrade() {
           ...(targetPlan === 'team' && { seats: CONSTANTS.INITIAL_TEAM_SEATS }),
         } as const
 
-        // Add subscriptionId for existing subscriptions to ensure proper plan switching
         const finalParams = currentSubscriptionId
           ? { ...upgradeParams, subscriptionId: currentSubscriptionId }
           : upgradeParams
 
         logger.info(
           currentSubscriptionId ? 'Upgrading existing subscription' : 'Creating new subscription',
-          {
-            targetPlan,
-            currentSubscriptionId,
-            referenceId,
-          }
+          { targetPlan, currentSubscriptionId, referenceId }
         )
 
         await betterAuthSubscription.upgrade(finalParams)
 
-        // If upgrading to team plan, ensure the subscription is transferred to the organization
         if (targetPlan === 'team' && currentSubscriptionId && referenceId !== userId) {
           try {
             logger.info('Transferring subscription to organization after upgrade', {
@@ -174,7 +125,6 @@ export function useSubscriptionUpgrade() {
                 organizationId: referenceId,
                 error: text,
               })
-              // We don't throw here because the upgrade itself succeeded
             } else {
               logger.info('Successfully transferred subscription to organization', {
                 subscriptionId: currentSubscriptionId,
@@ -186,21 +136,16 @@ export function useSubscriptionUpgrade() {
           }
         }
 
-        // For team plans, refresh organization data to ensure UI updates
         if (targetPlan === 'team') {
           try {
             await queryClient.invalidateQueries({ queryKey: organizationKeys.lists() })
             logger.info('Refreshed organization data after team upgrade')
           } catch (error) {
             logger.warn('Failed to refresh organization data after upgrade', error)
-            // Don't fail the entire upgrade if data refresh fails
           }
         }
 
-        logger.info('Subscription upgrade completed successfully', {
-          targetPlan,
-          referenceId,
-        })
+        logger.info('Subscription upgrade completed successfully', { targetPlan, referenceId })
       } catch (error) {
         logger.error('Failed to initiate subscription upgrade:', error)
 
diff --git a/apps/sim/lib/billing/organization.ts b/apps/sim/lib/billing/organization.ts
@@ -76,23 +76,18 @@ async function createOrganizationWithOwner(
   return newOrg.id
 }
 
-/**
- * Create organization for team/enterprise plan upgrade
- */
 export async function createOrganizationForTeamPlan(
   userId: string,
   userName?: string,
   userEmail?: string,
   organizationSlug?: string
 ): Promise<string> {
   try {
-    // Check if user already owns an organization
     const existingOrgId = await getUserOwnedOrganization(userId)
     if (existingOrgId) {
       return existingOrgId
     }
 
-    // Create new organization (same naming for both team and enterprise)
     const organizationName = userName || `${userEmail || 'User'}'s Team`
     const slug = organizationSlug || `${userId}-team-${Date.now()}`
 
@@ -117,6 +112,84 @@ export async function createOrganizationForTeamPlan(
   }
 }
 
+export async function ensureOrganizationForTeamSubscription(
+  subscription: SubscriptionData
+): Promise<SubscriptionData> {
+  if (subscription.plan !== 'team') {
+    return subscription
+  }
+
+  if (subscription.referenceId.startsWith('org_')) {
+    return subscription
+  }
+
+  const userId = subscription.referenceId
+
+  logger.info('Creating organization for team subscription', {
+    subscriptionId: subscription.id,
+    userId,
+  })
+
+  const existingMembership = await db
+    .select({
+      id: schema.member.id,
+      organizationId: schema.member.organizationId,
+      role: schema.member.role,
+    })
+    .from(schema.member)
+    .where(eq(schema.member.userId, userId))
+    .limit(1)
+
+  if (existingMembership.length > 0) {
+    const membership = existingMembership[0]
+    if (membership.role === 'owner' || membership.role === 'admin') {
+      logger.info('User already owns/admins an org, using it', {
+        userId,
+        organizationId: membership.organizationId,
+      })
+
+      await db
+        .update(schema.subscription)
+        .set({ referenceId: membership.organizationId })
+        .where(eq(schema.subscription.id, subscription.id))
+
+      return { ...subscription, referenceId: membership.organizationId }
+    }
+
+    logger.error('User is member of org but not owner/admin - cannot create team subscription', {
+      userId,
+      existingOrgId: membership.organizationId,
+      subscriptionId: subscription.id,
+    })
+    throw new Error('User is already member of another organization')
+  }
+
+  const [userData] = await db
+    .select({ name: schema.user.name, email: schema.user.email })
+    .from(schema.user)
+    .where(eq(schema.user.id, userId))
+    .limit(1)
+
+  const orgId = await createOrganizationForTeamPlan(
+    userId,
+    userData?.name || undefined,
+    userData?.email || undefined
+  )
+
+  await db
+    .update(schema.subscription)
+    .set({ referenceId: orgId })
+    .where(eq(schema.subscription.id, subscription.id))
+
+  logger.info('Created organization and updated subscription referenceId', {
+    subscriptionId: subscription.id,
+    userId,
+    organizationId: orgId,
+  })
+
+  return { ...subscription, referenceId: orgId }
+}
+
 /**
  * Sync usage limits for subscription members
  * Updates usage limits for all users associated with the subscription
