restore functionality fixed

Author: icecrasher321

apps/sim/app/workspace/[workspaceId]/w/components/sidebar/components/settings-modal/components/subscription/components/cancel-subscription/cancel-subscription.tsx

@@ -29,6 +29,7 @@ interface CancelSubscriptionProps {
   }
   subscriptionData?: {
     periodEnd?: Date | null
+    cancelAtPeriodEnd?: boolean
   }
 }
 
@@ -126,35 +127,48 @@ export function CancelSubscription({ subscription, subscriptionData }: CancelSub
       const subscriptionStatus = getSubscriptionStatus()
       const activeOrgId = activeOrganization?.id
 
-      // For team/enterprise plans, get the subscription ID from organization store
-      if ((subscriptionStatus.isTeam || subscriptionStatus.isEnterprise) && activeOrgId) {
-        const orgSubscription = useOrganizationStore.getState().subscriptionData
+      if (isCancelAtPeriodEnd) {
+        if (!betterAuthSubscription.restore) {
+          throw new Error('Subscription restore not available')
+        }
+
+        let referenceId: string
+        let subscriptionId: string | undefined
+
+        if ((subscriptionStatus.isTeam || subscriptionStatus.isEnterprise) && activeOrgId) {
+          const orgSubscription = useOrganizationStore.getState().subscriptionData
+          referenceId = activeOrgId
+          subscriptionId = orgSubscription?.id
+        } else {
+          // For personal subscriptions, use user ID and let better-auth find the subscription
+          referenceId = session.user.id
+          subscriptionId = undefined
+        }
+
+        logger.info('Restoring subscription', { referenceId, subscriptionId })
 
-        if (orgSubscription?.id && orgSubscription?.cancelAtPeriodEnd) {
-          // Restore the organization subscription
-          if (!betterAuthSubscription.restore) {
-            throw new Error('Subscription restore not available')
-          }
-
-          const result = await betterAuthSubscription.restore({
-            referenceId: activeOrgId,
-            subscriptionId: orgSubscription.id,
-          })
-          logger.info('Organization subscription restored successfully', result)
+        // Build restore params - only include subscriptionId if we have one (team/enterprise)
+        const restoreParams: any = { referenceId }
+        if (subscriptionId) {
+          restoreParams.subscriptionId = subscriptionId
         }
+
+        const result = await betterAuthSubscription.restore(restoreParams)
+
+        logger.info('Subscription restored successfully', result)
       }
 
-      // Refresh state and close
       await refresh()
       if (activeOrgId) {
         await loadOrganizationSubscription(activeOrgId)
         await refreshOrganization().catch(() => {})
       }
+
       setIsDialogOpen(false)
     } catch (error) {
-      const errorMessage = error instanceof Error ? error.message : 'Failed to keep subscription'
+      const errorMessage = error instanceof Error ? error.message : 'Failed to restore subscription'
       setError(errorMessage)
-      logger.error('Failed to keep subscription', { error })
+      logger.error('Failed to restore subscription', { error })
     } finally {
       setIsLoading(false)
     }
@@ -190,6 +204,12 @@ export function CancelSubscription({ subscription, subscriptionData }: CancelSub
 
   // Check if subscription is set to cancel at period end
   const isCancelAtPeriodEnd = (() => {
+    // First check if it's passed via props (works for all subscription types)
+    if (subscriptionData?.cancelAtPeriodEnd === true) {
+      return true
+    }
+
+    // Fall back to checking organization store for team subscriptions
     const subscriptionStatus = getSubscriptionStatus()
     if (subscriptionStatus.isTeam || subscriptionStatus.isEnterprise) {
       return useOrganizationStore.getState().subscriptionData?.cancelAtPeriodEnd === true
@@ -271,12 +291,7 @@ export function CancelSubscription({ subscription, subscriptionData }: CancelSub
 
             {(() => {
               const subscriptionStatus = getSubscriptionStatus()
-              if (
-                subscriptionStatus.isPaid &&
-                (activeOrganization?.id
-                  ? useOrganizationStore.getState().subscriptionData?.cancelAtPeriodEnd
-                  : false)
-              ) {
+              if (subscriptionStatus.isPaid && isCancelAtPeriodEnd) {
                 return (
                   <AlertDialogAction
                     onClick={handleKeep}

apps/sim/app/workspace/[workspaceId]/w/components/sidebar/components/settings-modal/components/subscription/subscription.tsx

@@ -523,6 +523,7 @@ export function Subscription({ onOpenChange }: SubscriptionProps) {
               }}
               subscriptionData={{
                 periodEnd: subscriptionData?.periodEnd || null,
+                cancelAtPeriodEnd: subscriptionData?.cancelAtPeriodEnd,
               }}
             />
           </div>

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

@@ -220,6 +220,7 @@ export async function getSimplifiedBillingSummary(
   metadata: any
   stripeSubscriptionId: string | null
   periodEnd: Date | string | null
+  cancelAtPeriodEnd?: boolean
   // Usage details
   usage: {
     current: number
@@ -318,6 +319,7 @@ export async function getSimplifiedBillingSummary(
         metadata: subscription.metadata || null,
         stripeSubscriptionId: subscription.stripeSubscriptionId || null,
         periodEnd: subscription.periodEnd || null,
+        cancelAtPeriodEnd: subscription.cancelAtPeriodEnd || undefined,
         // Usage details
         usage: {
           current: usageData.currentUsage,
@@ -393,6 +395,7 @@ export async function getSimplifiedBillingSummary(
       metadata: subscription?.metadata || null,
       stripeSubscriptionId: subscription?.stripeSubscriptionId || null,
       periodEnd: subscription?.periodEnd || null,
+      cancelAtPeriodEnd: subscription?.cancelAtPeriodEnd || undefined,
       // Usage details
       usage: {
         current: currentUsage,
@@ -450,5 +453,14 @@ function getDefaultBillingSummary(type: 'individual' | 'organization') {
       lastPeriodCost: 0,
       daysRemaining: 0,
     },
+    ...(type === 'organization' && {
+      organizationData: {
+        seatCount: 0,
+        memberCount: 0,
+        totalBasePrice: 0,
+        totalCurrentUsage: 0,
+        totalOverage: 0,
+      },
+    }),
   }
 }

apps/sim/stores/subscription/types.ts

@@ -29,6 +29,7 @@ export interface SubscriptionData {
   metadata: any | null
   stripeSubscriptionId: string | null
   periodEnd: Date | null
+  cancelAtPeriodEnd?: boolean
   usage: UsageData
   billingBlocked?: boolean
 }