fix(billing): reset usage on transition from free -> paid plan (#1397)

* fix(billing): reset usage on transition from free -> paid plan

* fixes: pro->team upgrade logic, single org server side check on invite routes

* ui improvements

* cleanup team-members code

* minor renaming

* progress

* fix pro->team upgrade to prevent double billing

* add subscription delete case handler

---------

Co-authored-by: Vikhyath Mondreti <[email protected]>

Author: waleedlatif1

apps/sim/app/api/organizations/[id]/invitations/[invitationId]/route.ts

@@ -5,13 +5,16 @@ import {
   member,
   organization,
   permissions,
+  subscription as subscriptionTable,
   user,
+  userStats,
   type WorkspaceInvitationStatus,
   workspaceInvitation,
 } from '@sim/db/schema'
 import { and, eq } from 'drizzle-orm'
 import { type NextRequest, NextResponse } from 'next/server'
 import { getSession } from '@/lib/auth'
+import { requireStripeClient } from '@/lib/billing/stripe-client'
 import { createLogger } from '@/lib/logs/console/logger'
 
 const logger = createLogger('OrganizationInvitation')
@@ -64,6 +67,16 @@ export async function PUT(
   { params }: { params: Promise<{ id: string; invitationId: string }> }
 ) {
   const { id: organizationId, invitationId } = await params
+
+  logger.info(
+    '[PUT /api/organizations/[id]/invitations/[invitationId]] Invitation acceptance request',
+    {
+      organizationId,
+      invitationId,
+      path: req.url,
+    }
+  )
+
   const session = await getSession()
 
   if (!session?.user?.id) {
@@ -130,6 +143,48 @@ export async function PUT(
       }
     }
 
+    // Enforce: user can only be part of a single organization
+    if (status === 'accepted') {
+      // Check if user is already a member of ANY organization
+      const existingOrgMemberships = await db
+        .select({ organizationId: member.organizationId })
+        .from(member)
+        .where(eq(member.userId, session.user.id))
+
+      if (existingOrgMemberships.length > 0) {
+        // Check if already a member of THIS specific organization
+        const alreadyMemberOfThisOrg = existingOrgMemberships.some(
+          (m) => m.organizationId === organizationId
+        )
+
+        if (alreadyMemberOfThisOrg) {
+          return NextResponse.json(
+            { error: 'You are already a member of this organization' },
+            { status: 400 }
+          )
+        }
+
+        // Member of a different organization
+        // Mark the invitation as rejected since they can't accept it
+        await db
+          .update(invitation)
+          .set({
+            status: 'rejected',
+          })
+          .where(eq(invitation.id, invitationId))
+
+        return NextResponse.json(
+          {
+            error:
+              'You are already a member of an organization. Leave your current organization before accepting a new invitation.',
+          },
+          { status: 409 }
+        )
+      }
+    }
+
+    let personalProToCancel: any = null
+
     await db.transaction(async (tx) => {
       await tx.update(invitation).set({ status }).where(eq(invitation.id, invitationId))
 
@@ -142,6 +197,83 @@ export async function PUT(
           createdAt: new Date(),
         })
 
+        // Snapshot Pro usage and cancel Pro subscription when joining a paid team
+        try {
+          const orgSubs = await tx
+            .select()
+            .from(subscriptionTable)
+            .where(
+              and(
+                eq(subscriptionTable.referenceId, organizationId),
+                eq(subscriptionTable.status, 'active')
+              )
+            )
+            .limit(1)
+
+          const orgSub = orgSubs[0]
+          const orgIsPaid = orgSub && (orgSub.plan === 'team' || orgSub.plan === 'enterprise')
+
+          if (orgIsPaid) {
+            const userId = session.user.id
+
+            // Find user's active personal Pro subscription
+            const personalSubs = await tx
+              .select()
+              .from(subscriptionTable)
+              .where(
+                and(
+                  eq(subscriptionTable.referenceId, userId),
+                  eq(subscriptionTable.status, 'active'),
+                  eq(subscriptionTable.plan, 'pro')
+                )
+              )
+              .limit(1)
+
+            const personalPro = personalSubs[0]
+            if (personalPro) {
+              // Snapshot the current Pro usage before resetting
+              const userStatsRows = await tx
+                .select({
+                  currentPeriodCost: userStats.currentPeriodCost,
+                })
+                .from(userStats)
+                .where(eq(userStats.userId, userId))
+                .limit(1)
+
+              if (userStatsRows.length > 0) {
+                const currentProUsage = userStatsRows[0].currentPeriodCost || '0'
+
+                // Snapshot Pro usage and reset currentPeriodCost so new usage goes to team
+                await tx
+                  .update(userStats)
+                  .set({
+                    proPeriodCostSnapshot: currentProUsage,
+                    currentPeriodCost: '0', // Reset so new usage is attributed to team
+                  })
+                  .where(eq(userStats.userId, userId))
+
+                logger.info('Snapshotted Pro usage when joining team', {
+                  userId,
+                  proUsageSnapshot: currentProUsage,
+                  organizationId,
+                })
+              }
+
+              // Mark for cancellation after transaction
+              if (personalPro.cancelAtPeriodEnd !== true) {
+                personalProToCancel = personalPro
+              }
+            }
+          }
+        } catch (error) {
+          logger.error('Failed to handle Pro user joining team', {
+            userId: session.user.id,
+            organizationId,
+            error,
+          })
+          // Don't fail the whole invitation acceptance due to this
+        }
+
         const linkedWorkspaceInvitations = await tx
           .select()
           .from(workspaceInvitation)
@@ -179,6 +311,44 @@ export async function PUT(
       }
     })
 
+    // Handle Pro subscription cancellation after transaction commits
+    if (personalProToCancel) {
+      try {
+        const stripe = requireStripeClient()
+        if (personalProToCancel.stripeSubscriptionId) {
+          try {
+            await stripe.subscriptions.update(personalProToCancel.stripeSubscriptionId, {
+              cancel_at_period_end: true,
+            })
+          } catch (stripeError) {
+            logger.error('Failed to set cancel_at_period_end on Stripe for personal Pro', {
+              userId: session.user.id,
+              subscriptionId: personalProToCancel.id,
+              stripeSubscriptionId: personalProToCancel.stripeSubscriptionId,
+              error: stripeError,
+            })
+          }
+        }
+
+        await db
+          .update(subscriptionTable)
+          .set({ cancelAtPeriodEnd: true })
+          .where(eq(subscriptionTable.id, personalProToCancel.id))
+
+        logger.info('Auto-cancelled personal Pro at period end after joining paid team', {
+          userId: session.user.id,
+          personalSubscriptionId: personalProToCancel.id,
+          organizationId,
+        })
+      } catch (dbError) {
+        logger.error('Failed to update DB cancelAtPeriodEnd for personal Pro', {
+          userId: session.user.id,
+          subscriptionId: personalProToCancel.id,
+          error: dbError,
+        })
+      }
+    }
+
     logger.info(`Organization invitation ${status}`, {
       organizationId,
       invitationId,

apps/sim/app/api/organizations/[id]/invitations/route.ts

@@ -9,7 +9,7 @@ import {
   workspace,
   workspaceInvitation,
 } from '@sim/db/schema'
-import { and, eq, inArray, isNull } from 'drizzle-orm'
+import { and, eq, inArray, isNull, or } from 'drizzle-orm'
 import { type NextRequest, NextResponse } from 'next/server'
 import {
   getEmailSubject,
@@ -463,7 +463,10 @@ export async function DELETE(
         and(
           eq(invitation.id, invitationId),
           eq(invitation.organizationId, organizationId),
-          eq(invitation.status, 'pending')
+          or(
+            eq(invitation.status, 'pending'),
+            eq(invitation.status, 'rejected') // Allow cancelling rejected invitations too
+          )
         )
       )
       .returning()

apps/sim/app/api/organizations/[id]/members/[memberId]/route.ts

@@ -1,9 +1,10 @@
 import { db } from '@sim/db'
-import { member, user, userStats } from '@sim/db/schema'
+import { member, subscription as subscriptionTable, user, userStats } from '@sim/db/schema'
 import { and, eq } from 'drizzle-orm'
 import { type NextRequest, NextResponse } from 'next/server'
 import { getSession } from '@/lib/auth'
 import { getUserUsageData } from '@/lib/billing/core/usage'
+import { requireStripeClient } from '@/lib/billing/stripe-client'
 import { createLogger } from '@/lib/logs/console/logger'
 
 const logger = createLogger('OrganizationMemberAPI')
@@ -304,6 +305,124 @@ export async function DELETE(
       wasSelfRemoval: session.user.id === memberId,
     })
 
+    // If the removed user left their last paid team and has a personal Pro set to cancel_at_period_end, restore it
+    try {
+      const remainingPaidTeams = await db
+        .select({ orgId: member.organizationId })
+        .from(member)
+        .where(eq(member.userId, memberId))
+
+      let hasAnyPaidTeam = false
+      if (remainingPaidTeams.length > 0) {
+        const orgIds = remainingPaidTeams.map((m) => m.orgId)
+        const orgPaidSubs = await db
+          .select()
+          .from(subscriptionTable)
+          .where(and(eq(subscriptionTable.status, 'active'), eq(subscriptionTable.plan, 'team')))
+
+        hasAnyPaidTeam = orgPaidSubs.some((s) => orgIds.includes(s.referenceId))
+      }
+
+      if (!hasAnyPaidTeam) {
+        const personalProRows = await db
+          .select()
+          .from(subscriptionTable)
+          .where(
+            and(
+              eq(subscriptionTable.referenceId, memberId),
+              eq(subscriptionTable.status, 'active'),
+              eq(subscriptionTable.plan, 'pro')
+            )
+          )
+          .limit(1)
+
+        const personalPro = personalProRows[0]
+        if (
+          personalPro &&
+          personalPro.cancelAtPeriodEnd === true &&
+          personalPro.stripeSubscriptionId
+        ) {
+          try {
+            const stripe = requireStripeClient()
+            await stripe.subscriptions.update(personalPro.stripeSubscriptionId, {
+              cancel_at_period_end: false,
+            })
+          } catch (stripeError) {
+            logger.error('Stripe restore cancel_at_period_end failed for personal Pro', {
+              userId: memberId,
+              stripeSubscriptionId: personalPro.stripeSubscriptionId,
+              error: stripeError,
+            })
+          }
+
+          try {
+            await db
+              .update(subscriptionTable)
+              .set({ cancelAtPeriodEnd: false })
+              .where(eq(subscriptionTable.id, personalPro.id))
+
+            logger.info('Restored personal Pro after leaving last paid team', {
+              userId: memberId,
+              personalSubscriptionId: personalPro.id,
+            })
+          } catch (dbError) {
+            logger.error('DB update failed when restoring personal Pro', {
+              userId: memberId,
+              subscriptionId: personalPro.id,
+              error: dbError,
+            })
+          }
+
+          // Also restore the snapshotted Pro usage back to currentPeriodCost
+          try {
+            const userStatsRows = await db
+              .select({
+                currentPeriodCost: userStats.currentPeriodCost,
+                proPeriodCostSnapshot: userStats.proPeriodCostSnapshot,
+              })
+              .from(userStats)
+              .where(eq(userStats.userId, memberId))
+              .limit(1)
+
+            if (userStatsRows.length > 0) {
+              const currentUsage = userStatsRows[0].currentPeriodCost || '0'
+              const snapshotUsage = userStatsRows[0].proPeriodCostSnapshot || '0'
+
+              const currentNum = Number.parseFloat(currentUsage)
+              const snapshotNum = Number.parseFloat(snapshotUsage)
+              const restoredUsage = (currentNum + snapshotNum).toString()
+
+              await db
+                .update(userStats)
+                .set({
+                  currentPeriodCost: restoredUsage,
+                  proPeriodCostSnapshot: '0', // Clear the snapshot
+                })
+                .where(eq(userStats.userId, memberId))
+
+              logger.info('Restored Pro usage after leaving team', {
+                userId: memberId,
+                previousUsage: currentUsage,
+                snapshotUsage: snapshotUsage,
+                restoredUsage: restoredUsage,
+              })
+            }
+          } catch (usageRestoreError) {
+            logger.error('Failed to restore Pro usage after leaving team', {
+              userId: memberId,
+              error: usageRestoreError,
+            })
+          }
+        }
+      }
+    } catch (postRemoveError) {
+      logger.error('Post-removal personal Pro restore check failed', {
+        organizationId,
+        memberId,
+        error: postRemoveError,
+      })
+    }
+
     return NextResponse.json({
       success: true,
       message: