fix(billing): fix team plan upgrade (#1053)

apps/sim/lib/auth.ts

@@ -1270,133 +1270,30 @@ export const auth = betterAuth({
                 })
 
                 // Auto-create organization for team plan purchases
-                if (subscription.plan === 'team') {
-                  try {
-                    // Get the user who purchased the subscription
-                    const user = await db
-                      .select()
-                      .from(schema.user)
-                      .where(eq(schema.user.id, subscription.referenceId))
-                      .limit(1)
-
-                    if (user.length > 0) {
-                      const currentUser = user[0]
-
-                      // Store the original user ID before we change the referenceId
-                      const originalUserId = subscription.referenceId
-
-                      // First, sync usage limits for the purchasing user with their new plan
-                      try {
-                        const { syncUsageLimitsFromSubscription } = await import('@/lib/billing')
-                        await syncUsageLimitsFromSubscription(originalUserId)
-                        logger.info(
-                          'Usage limits synced for purchasing user before organization creation',
-                          {
-                            userId: originalUserId,
-                          }
-                        )
-                      } catch (error) {
-                        logger.error('Failed to sync usage limits for purchasing user', {
-                          userId: originalUserId,
-                          error,
-                        })
-                      }
-
-                      // Create organization for the team
-                      const orgId = `org_${Date.now()}_${Math.random().toString(36).substring(2, 15)}`
-                      const orgSlug = `${currentUser.name?.toLowerCase().replace(/\s+/g, '-') || 'team'}-${Date.now()}`
-
-                      // Create a separate Stripe customer for the organization
-                      let orgStripeCustomerId: string | null = null
-                      if (stripeClient) {
-                        try {
-                          const orgStripeCustomer = await stripeClient.customers.create({
-                            name: `${currentUser.name || 'User'}'s Team`,
-                            email: currentUser.email,
-                            metadata: {
-                              organizationId: orgId,
-                              type: 'organization',
-                            },
-                          })
-                          orgStripeCustomerId = orgStripeCustomer.id
-                        } catch (error) {
-                          logger.error('Failed to create Stripe customer for organization', {
-                            organizationId: orgId,
-                            error,
-                          })
-                          // Continue without Stripe customer - can be created later
-                        }
-                      }
-
-                      const newOrg = await db
-                        .insert(schema.organization)
-                        .values({
-                          id: orgId,
-                          name: `${currentUser.name || 'User'}'s Team`,
-                          slug: orgSlug,
-                          metadata: orgStripeCustomerId
-                            ? { stripeCustomerId: orgStripeCustomerId }
-                            : null,
-                          createdAt: new Date(),
-                          updatedAt: new Date(),
-                        })
-                        .returning()
-
-                      // Add the user as owner of the organization
-                      await db.insert(schema.member).values({
-                        id: `member_${Date.now()}_${Math.random().toString(36).substring(2, 15)}`,
-                        userId: currentUser.id,
-                        organizationId: orgId,
-                        role: 'owner',
-                        createdAt: new Date(),
-                      })
-
-                      // Update the subscription to reference the organization instead of the user
-                      await db
-                        .update(schema.subscription)
-                        .set({ referenceId: orgId })
-                        .where(eq(schema.subscription.id, subscription.id))
-
-                      // Update the session to set the new organization as active
-                      await db
-                        .update(schema.session)
-                        .set({ activeOrganizationId: orgId })
-                        .where(eq(schema.session.userId, currentUser.id))
-
-                      logger.info('Auto-created organization for team subscription', {
-                        organizationId: orgId,
-                        userId: currentUser.id,
-                        subscriptionId: subscription.id,
-                        orgName: `${currentUser.name || 'User'}'s Team`,
-                      })
-
-                      // Update referenceId for usage limit sync
-                      subscription.referenceId = orgId
-                    }
-                  } catch (error) {
-                    logger.error('Failed to auto-create organization for team subscription', {
-                      subscriptionId: subscription.id,
-                      referenceId: subscription.referenceId,
-                      error,
-                    })
-                  }
+                try {
+                  const { handleTeamPlanOrganization } = await import(
+                    '@/lib/billing/team-management'
+                  )
+                  await handleTeamPlanOrganization(subscription)
+                } catch (error) {
+                  logger.error('Failed to handle team plan organization creation', {
+                    subscriptionId: subscription.id,
+                    referenceId: subscription.referenceId,
+                    error,
+                  })
                 }
 
-                // Initialize billing period for the user/organization
+                // Initialize billing period and sync usage limits
                 try {
                   const { initializeBillingPeriod } = await import(
                     '@/lib/billing/core/billing-periods'
                   )
+                  const { syncSubscriptionUsageLimits } = await import(
+                    '@/lib/billing/team-management'
+                  )
 
-                  // Note: Usage limits are already synced above for team plan users
-                  // For non-team plans, sync usage limits using the referenceId (which is the user ID)
-                  if (subscription.plan !== 'team') {
-                    const { syncUsageLimitsFromSubscription } = await import('@/lib/billing')
-                    await syncUsageLimitsFromSubscription(subscription.referenceId)
-                    logger.info('Usage limits synced after subscription creation', {
-                      referenceId: subscription.referenceId,
-                    })
-                  }
+                  // Sync usage limits for user or organization members
+                  await syncSubscriptionUsageLimits(subscription)
 
                   // Initialize billing period for new subscription using Stripe dates
                   if (subscription.plan !== 'free') {
@@ -1433,15 +1330,29 @@ export const auth = betterAuth({
                 logger.info('Subscription updated', {
                   subscriptionId: subscription.id,
                   status: subscription.status,
+                  plan: subscription.plan,
                 })
 
+                // Auto-create organization for team plan upgrades (free -> team)
+                try {
+                  const { handleTeamPlanOrganization } = await import(
+                    '@/lib/billing/team-management'
+                  )
+                  await handleTeamPlanOrganization(subscription)
+                } catch (error) {
+                  logger.error('Failed to handle team plan organization creation on update', {
+                    subscriptionId: subscription.id,
+                    referenceId: subscription.referenceId,
+                    error,
+                  })
+                }
+
                 // Sync usage limits for the user/organization
                 try {
-                  const { syncUsageLimitsFromSubscription } = await import('@/lib/billing')
-                  await syncUsageLimitsFromSubscription(subscription.referenceId)
-                  logger.info('Usage limits synced after subscription update', {
-                    referenceId: subscription.referenceId,
-                  })
+                  const { syncSubscriptionUsageLimits } = await import(
+                    '@/lib/billing/team-management'
+                  )
+                  await syncSubscriptionUsageLimits(subscription)
                 } catch (error) {
                   logger.error('Failed to sync usage limits after subscription update', {
                     referenceId: subscription.referenceId,

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

@@ -13,6 +13,24 @@ import { member, organization, subscription, user, userStats } from '@/db/schema
 
 const logger = createLogger('Billing')
 
+/**
+ * Get organization subscription directly by organization ID
+ */
+export async function getOrganizationSubscription(organizationId: string) {
+  try {
+    const orgSubs = await db
+      .select()
+      .from(subscription)
+      .where(and(eq(subscription.referenceId, organizationId), eq(subscription.status, 'active')))
+      .limit(1)
+
+    return orgSubs.length > 0 ? orgSubs[0] : null
+  } catch (error) {
+    logger.error('Error getting organization subscription', { error, organizationId })
+    return null
+  }
+}
+
 interface BillingResult {
   success: boolean
   chargedAmount?: number
@@ -89,15 +107,43 @@ async function getStripeCustomerId(referenceId: string): Promise<string | null>
       .where(eq(organization.id, referenceId))
       .limit(1)
 
-    if (orgRecord.length > 0 && orgRecord[0].metadata) {
-      const metadata =
-        typeof orgRecord[0].metadata === 'string'
-          ? JSON.parse(orgRecord[0].metadata)
-          : orgRecord[0].metadata
+    if (orgRecord.length > 0) {
+      // First, check if organization has its own Stripe customer (legacy support)
+      if (orgRecord[0].metadata) {
+        const metadata =
+          typeof orgRecord[0].metadata === 'string'
+            ? JSON.parse(orgRecord[0].metadata)
+            : orgRecord[0].metadata
 
-      if (metadata?.stripeCustomerId) {
-        return metadata.stripeCustomerId
+        if (metadata?.stripeCustomerId) {
+          return metadata.stripeCustomerId
+        }
       }
+
+      // If organization has no Stripe customer, use the owner's customer
+      // This is our new pattern: subscriptions stay with user, referenceId = orgId
+      const ownerRecord = await db
+        .select({
+          stripeCustomerId: user.stripeCustomerId,
+          userId: user.id,
+        })
+        .from(user)
+        .innerJoin(member, eq(member.userId, user.id))
+        .where(and(eq(member.organizationId, referenceId), eq(member.role, 'owner')))
+        .limit(1)
+
+      if (ownerRecord.length > 0 && ownerRecord[0].stripeCustomerId) {
+        logger.debug('Using organization owner Stripe customer for billing', {
+          organizationId: referenceId,
+          ownerId: ownerRecord[0].userId,
+          stripeCustomerId: ownerRecord[0].stripeCustomerId,
+        })
+        return ownerRecord[0].stripeCustomerId
+      }
+
+      logger.warn('No Stripe customer found for organization or its owner', {
+        organizationId: referenceId,
+      })
     }
 
     return null
@@ -431,8 +477,8 @@ export async function processOrganizationOverageBilling(
   organizationId: string
 ): Promise<BillingResult> {
   try {
-    // Get organization subscription
-    const subscription = await getHighestPrioritySubscription(organizationId)
+    // Get organization subscription directly (referenceId = organizationId)
+    const subscription = await getOrganizationSubscription(organizationId)
 
     if (!subscription || !['team', 'enterprise'].includes(subscription.plan)) {
       logger.warn('No team/enterprise subscription found for organization', { organizationId })
@@ -759,7 +805,9 @@ export async function getSimplifiedBillingSummary(
   try {
     // Get subscription and usage data upfront
     const [subscription, usageData] = await Promise.all([
-      getHighestPrioritySubscription(organizationId || userId),
+      organizationId
+        ? getOrganizationSubscription(organizationId)
+        : getHighestPrioritySubscription(userId),
       getUserUsageData(userId),
     ])
 

apps/sim/lib/billing/core/organization-billing.ts

@@ -1,13 +1,31 @@
 import { and, eq } from 'drizzle-orm'
 import { DEFAULT_FREE_CREDITS } from '@/lib/billing/constants'
 import { getPlanPricing } from '@/lib/billing/core/billing'
-import { getHighestPrioritySubscription } from '@/lib/billing/core/subscription'
 import { createLogger } from '@/lib/logs/console/logger'
 import { db } from '@/db'
-import { member, organization, user, userStats } from '@/db/schema'
+import { member, organization, subscription, user, userStats } from '@/db/schema'
 
 const logger = createLogger('OrganizationBilling')
 
+/**
+ * Get organization subscription directly by organization ID
+ * This is for our new pattern where referenceId = organizationId
+ */
+async function getOrganizationSubscription(organizationId: string) {
+  try {
+    const orgSubs = await db
+      .select()
+      .from(subscription)
+      .where(and(eq(subscription.referenceId, organizationId), eq(subscription.status, 'active')))
+      .limit(1)
+
+    return orgSubs.length > 0 ? orgSubs[0] : null
+  } catch (error) {
+    logger.error('Error getting organization subscription', { error, organizationId })
+    return null
+  }
+}
+
 interface OrganizationUsageData {
   organizationId: string
   organizationName: string
@@ -57,8 +75,8 @@ export async function getOrganizationBillingData(
 
     const organizationData = orgRecord[0]
 
-    // Get organization subscription
-    const subscription = await getHighestPrioritySubscription(organizationId)
+    // Get organization subscription directly (referenceId = organizationId)
+    const subscription = await getOrganizationSubscription(organizationId)
 
     if (!subscription) {
       logger.warn('No subscription found for organization', { organizationId })
@@ -191,7 +209,7 @@ export async function updateMemberUsageLimit(
     }
 
     // Get organization subscription to validate limit
-    const subscription = await getHighestPrioritySubscription(organizationId)
+    const subscription = await getOrganizationSubscription(organizationId)
     if (!subscription) {
       throw new Error('No active subscription found')
     }

apps/sim/lib/billing/team-management.ts

@@ -0,0 +1,181 @@
+import { eq } from 'drizzle-orm'
+import { createLogger } from '@/lib/logs/console/logger'
+import { db } from '@/db'
+import { member, organization, session, subscription, user } from '@/db/schema'
+
+const logger = createLogger('TeamManagement')
+
+type SubscriptionData = {
+  id: string
+  plan: string
+  referenceId: string
+  status: string
+  seats?: number
+  [key: string]: any
+}
+
+/**
+ * Auto-create organization for team plan subscriptions
+ */
+export async function handleTeamPlanOrganization(
+  subscriptionData: SubscriptionData
+): Promise<void> {
+  if (subscriptionData.plan !== 'team') return
+
+  try {
+    // For team subscriptions, referenceId should be the user ID initially
+    // But if the organization has already been created, it might be the org ID
+    let userId: string = subscriptionData.referenceId
+    let currentUser: any = null
+
+    // First try to get user directly (most common case)
+    const users = await db
+      .select()
+      .from(user)
+      .where(eq(user.id, subscriptionData.referenceId))
+      .limit(1)
+
+    if (users.length > 0) {
+      currentUser = users[0]
+      userId = currentUser.id
+    } else {
+      // If referenceId is not a user ID, it might be an organization ID
+      // In that case, the organization already exists, so we should skip
+      const existingOrg = await db
+        .select()
+        .from(organization)
+        .where(eq(organization.id, subscriptionData.referenceId))
+        .limit(1)
+
+      if (existingOrg.length > 0) {
+        logger.info('Organization already exists for team subscription, skipping creation', {
+          organizationId: subscriptionData.referenceId,
+          subscriptionId: subscriptionData.id,
+        })
+        return
+      }
+
+      logger.warn('User not found for team subscription and no existing organization', {
+        referenceId: subscriptionData.referenceId,
+      })
+      return
+    }
+
+    // Check if user already has an organization membership
+    const existingMember = await db.select().from(member).where(eq(member.userId, userId)).limit(1)
+
+    if (existingMember.length > 0) {
+      logger.info('User already has organization membership, skipping auto-creation', {
+        userId,
+        existingOrgId: existingMember[0].organizationId,
+      })
+      return
+    }
+
+    const orgName = `${currentUser.name || 'User'}'s Team`
+    const orgSlug = `${currentUser.name?.toLowerCase().replace(/\s+/g, '-') || 'team'}-${Date.now()}`
+
+    // Create organization directly in database
+    const orgId = `org_${Date.now()}_${Math.random().toString(36).substring(2, 15)}`
+
+    const [createdOrg] = await db
+      .insert(organization)
+      .values({
+        id: orgId,
+        name: orgName,
+        slug: orgSlug,
+        createdAt: new Date(),
+        updatedAt: new Date(),
+      })
+      .returning()
+
+    if (!createdOrg) {
+      throw new Error('Failed to create organization in database')
+    }
+
+    // Add the user as admin of the organization (owner role for full control)
+    await db.insert(member).values({
+      id: `member_${Date.now()}_${Math.random().toString(36).substring(2, 15)}`,
+      userId: currentUser.id,
+      organizationId: orgId,
+      role: 'owner', // Owner gives full admin privileges
+      createdAt: new Date(),
+    })
+
+    // Update the subscription to reference the organization instead of the user
+    await db
+      .update(subscription)
+      .set({ referenceId: orgId })
+      .where(eq(subscription.id, subscriptionData.id))
+
+    // Update the user's session to set the new organization as active
+    await db
+      .update(session)
+      .set({ activeOrganizationId: orgId })
+      .where(eq(session.userId, currentUser.id))
+
+    logger.info('Auto-created organization for team subscription', {
+      organizationId: orgId,
+      userId: currentUser.id,
+      subscriptionId: subscriptionData.id,
+      orgName,
+      userRole: 'owner',
+    })
+
+    // Update subscription object for subsequent logic
+    subscriptionData.referenceId = orgId
+  } catch (error) {
+    logger.error('Failed to auto-create organization for team subscription', {
+      subscriptionId: subscriptionData.id,
+      referenceId: subscriptionData.referenceId,
+      error,
+    })
+    throw error
+  }
+}
+
+/**
+ * Sync usage limits for user or organization
+ * Handles the complexity of determining whether to sync for user ID or org members
+ */
+export async function syncSubscriptionUsageLimits(
+  subscriptionData: SubscriptionData
+): Promise<void> {
+  try {
+    const { syncUsageLimitsFromSubscription } = await import('@/lib/billing')
+
+    // For team plans, the referenceId is now an organization ID
+    // We need to sync limits for the organization members
+    if (subscriptionData.plan === 'team') {
+      // Get all members of the organization
+      const orgMembers = await db
+        .select({ userId: member.userId })
+        .from(member)
+        .where(eq(member.organizationId, subscriptionData.referenceId))
+
+      // Sync usage limits for each member
+      for (const orgMember of orgMembers) {
+        await syncUsageLimitsFromSubscription(orgMember.userId)
+      }
+
+      logger.info('Synced usage limits for team organization members', {
+        organizationId: subscriptionData.referenceId,
+        memberCount: orgMembers.length,
+      })
+    } else {
+      // For non-team plans, referenceId is the user ID
+      await syncUsageLimitsFromSubscription(subscriptionData.referenceId)
+      logger.info('Synced usage limits for user', {
+        userId: subscriptionData.referenceId,
+        plan: subscriptionData.plan,
+      })
+    }
+  } catch (error) {
+    logger.error('Failed to sync subscription usage limits', {
+      subscriptionId: subscriptionData.id,
+      referenceId: subscriptionData.referenceId,
+      error,
+    })
+    throw error
+  }
+}

apps/sim/lib/billing/validation/seat-management.ts

@@ -1,5 +1,5 @@
 import { and, count, eq } from 'drizzle-orm'
-import { getHighestPrioritySubscription } from '@/lib/billing/core/subscription'
+import { getOrganizationSubscription } from '@/lib/billing/core/billing'
 import { quickValidateEmail } from '@/lib/email/validation'
 import { createLogger } from '@/lib/logs/console/logger'
 import { db } from '@/db'
@@ -33,8 +33,8 @@ export async function validateSeatAvailability(
   additionalSeats = 1
 ): Promise<SeatValidationResult> {
   try {
-    // Get organization subscription
-    const subscription = await getHighestPrioritySubscription(organizationId)
+    // Get organization subscription directly (referenceId = organizationId)
+    const subscription = await getOrganizationSubscription(organizationId)
 
     if (!subscription) {
       return {
@@ -71,7 +71,10 @@ export async function validateSeatAvailability(
     // For enterprise plans, check metadata for custom seat allowances
     if (subscription.plan === 'enterprise' && subscription.metadata) {
       try {
-        const metadata = JSON.parse(subscription.metadata)
+        const metadata =
+          typeof subscription.metadata === 'string'
+            ? JSON.parse(subscription.metadata)
+            : subscription.metadata
         if (metadata.maxSeats) {
           maxSeats = metadata.maxSeats
         }
@@ -142,8 +145,8 @@ export async function getOrganizationSeatInfo(
       return null
     }
 
-    // Get subscription
-    const subscription = await getHighestPrioritySubscription(organizationId)
+    // Get organization subscription directly (referenceId = organizationId)
+    const subscription = await getOrganizationSubscription(organizationId)
 
     if (!subscription) {
       return null
@@ -163,7 +166,10 @@ export async function getOrganizationSeatInfo(
 
     if (subscription.plan === 'enterprise' && subscription.metadata) {
       try {
-        const metadata = JSON.parse(subscription.metadata)
+        const metadata =
+          typeof subscription.metadata === 'string'
+            ? JSON.parse(subscription.metadata)
+            : subscription.metadata
         if (metadata.maxSeats) {
           maxSeats = metadata.maxSeats
         }
@@ -282,8 +288,8 @@ export async function updateOrganizationSeats(
   updatedBy: string
 ): Promise<{ success: boolean; error?: string }> {
   try {
-    // Get current subscription
-    const subscriptionRecord = await getHighestPrioritySubscription(organizationId)
+    // Get current organization subscription directly (referenceId = organizationId)
+    const subscriptionRecord = await getOrganizationSubscription(organizationId)
 
     if (!subscriptionRecord) {
       return { success: false, error: 'No active subscription found' }