improvement(billing): team upgrade + session management (#2751)

icecrasher321 · web-flow · commit 860610b4c2ad · 2026-01-09T16:36:45.000-08:00
* improvement(billng): team upgrade + session management

* remove comments

* session updates should be atomic

* make consistent for onSubscritionUpdate

* plan upgrade to refresh session

* fix var name

* remove dead code

* preserve params
diff --git a/apps/sim/app/_shell/providers/session-provider.tsx b/apps/sim/app/_shell/providers/session-provider.tsx
@@ -2,6 +2,7 @@
 
 import type React from 'react'
 import { createContext, useCallback, useEffect, useMemo, useState } from 'react'
+import { useQueryClient } from '@tanstack/react-query'
 import posthog from 'posthog-js'
 import { client } from '@/lib/auth/auth-client'
 
@@ -35,12 +36,15 @@ export function SessionProvider({ children }: { children: React.ReactNode }) {
   const [data, setData] = useState<AppSession>(null)
   const [isPending, setIsPending] = useState(true)
   const [error, setError] = useState<Error | null>(null)
+  const queryClient = useQueryClient()
 
-  const loadSession = useCallback(async () => {
+  const loadSession = useCallback(async (bypassCache = false) => {
     try {
       setIsPending(true)
       setError(null)
-      const res = await client.getSession()
+      const res = bypassCache
+        ? await client.getSession({ query: { disableCookieCache: true } })
+        : await client.getSession()
       setData(res?.data ?? null)
     } catch (e) {
       setError(e instanceof Error ? e : new Error('Failed to fetch session'))
@@ -50,8 +54,25 @@ export function SessionProvider({ children }: { children: React.ReactNode }) {
   }, [])
 
   useEffect(() => {
-    loadSession()
-  }, [loadSession])
+    // Check if user was redirected after plan upgrade
+    const params = new URLSearchParams(window.location.search)
+    const wasUpgraded = params.get('upgraded') === 'true'
+
+    if (wasUpgraded) {
+      params.delete('upgraded')
+      const newUrl = params.toString()
+        ? `${window.location.pathname}?${params.toString()}`
+        : window.location.pathname
+      window.history.replaceState({}, '', newUrl)
+    }
+
+    loadSession(wasUpgraded).then(() => {
+      if (wasUpgraded) {
+        queryClient.invalidateQueries({ queryKey: ['organizations'] })
+        queryClient.invalidateQueries({ queryKey: ['subscription'] })
+      }
+    })
+  }, [loadSession, queryClient])
 
   useEffect(() => {
     if (isPending || typeof posthog.identify !== 'function') {
diff --git a/apps/sim/app/workspace/[workspaceId]/w/components/sidebar/components/settings-modal/components/subscription/subscription.tsx b/apps/sim/app/workspace/[workspaceId]/w/components/sidebar/components/settings-modal/components/subscription/subscription.tsx
@@ -8,6 +8,7 @@ import { Skeleton } from '@/components/ui'
 import { useSession } from '@/lib/auth/auth-client'
 import { useSubscriptionUpgrade } from '@/lib/billing/client/upgrade'
 import { USAGE_THRESHOLDS } from '@/lib/billing/client/usage-visualization'
+import { getEffectiveSeats } from '@/lib/billing/subscriptions/utils'
 import { cn } from '@/lib/core/utils/cn'
 import { getBaseUrl } from '@/lib/core/utils/urls'
 import { getUserRole } from '@/lib/workspaces/organization/utils'
@@ -191,7 +192,13 @@ export function Subscription() {
   const [upgradeError, setUpgradeError] = useState<'pro' | 'team' | null>(null)
   const usageLimitRef = useRef<UsageLimitRef | null>(null)
 
-  const isLoading = isSubscriptionLoading || isUsageLimitLoading || isWorkspaceLoading
+  const isOrgPlan =
+    subscriptionData?.data?.plan === 'team' || subscriptionData?.data?.plan === 'enterprise'
+  const isLoading =
+    isSubscriptionLoading ||
+    isUsageLimitLoading ||
+    isWorkspaceLoading ||
+    (isOrgPlan && isOrgBillingLoading)
 
   const subscription = {
     isFree: subscriptionData?.data?.plan === 'free' || !subscriptionData?.data?.plan,
@@ -204,7 +211,7 @@ export function Subscription() {
       subscriptionData?.data?.status === 'active',
     plan: subscriptionData?.data?.plan || 'free',
     status: subscriptionData?.data?.status || 'inactive',
-    seats: organizationBillingData?.totalSeats ?? 0,
+    seats: getEffectiveSeats(subscriptionData?.data),
   }
 
   const usage = {
@@ -445,16 +452,10 @@ export function Subscription() {
             ? `${subscription.seats} seats`
             : undefined
         }
-        current={
-          subscription.isEnterprise || subscription.isTeam
-            ? (organizationBillingData?.totalCurrentUsage ?? usage.current)
-            : usage.current
-        }
+        current={usage.current}
         limit={
           subscription.isEnterprise || subscription.isTeam
-            ? organizationBillingData?.totalUsageLimit ||
-              organizationBillingData?.minimumBillingAmount ||
-              usage.limit
+            ? organizationBillingData?.data?.totalUsageLimit
             : !subscription.isFree &&
                 (permissions.canEditUsageLimit || permissions.showTeamMemberView)
               ? usage.current // placeholder; rightContent will render UsageLimit
@@ -468,19 +469,31 @@ export function Subscription() {
             <UsageLimit
               ref={usageLimitRef}
               currentLimit={
-                subscription.isTeam && isTeamAdmin
-                  ? organizationBillingData?.totalUsageLimit || usage.limit
+                (subscription.isTeam || subscription.isEnterprise) &&
+                isTeamAdmin &&
+                organizationBillingData?.data
+                  ? organizationBillingData.data.totalUsageLimit
                   : usageLimitData.currentLimit || usage.limit
               }
               currentUsage={usage.current}
               canEdit={permissions.canEditUsageLimit}
               minimumLimit={
-                subscription.isTeam && isTeamAdmin
-                  ? organizationBillingData?.minimumBillingAmount || (subscription.isPro ? 20 : 40)
+                (subscription.isTeam || subscription.isEnterprise) &&
+                isTeamAdmin &&
+                organizationBillingData?.data
+                  ? organizationBillingData.data.minimumBillingAmount
                   : usageLimitData.minimumLimit || (subscription.isPro ? 20 : 40)
               }
-              context={subscription.isTeam && isTeamAdmin ? 'organization' : 'user'}
-              organizationId={subscription.isTeam && isTeamAdmin ? activeOrgId : undefined}
+              context={
+                (subscription.isTeam || subscription.isEnterprise) && isTeamAdmin
+                  ? 'organization'
+                  : 'user'
+              }
+              organizationId={
+                (subscription.isTeam || subscription.isEnterprise) && isTeamAdmin
+                  ? activeOrgId
+                  : undefined
+              }
               onLimitUpdated={() => {
                 logger.info('Usage limit updated')
               }}
diff --git a/apps/sim/app/workspace/[workspaceId]/w/components/sidebar/components/settings-modal/settings-modal.tsx b/apps/sim/app/workspace/[workspaceId]/w/components/sidebar/components/settings-modal/settings-modal.tsx
@@ -174,6 +174,7 @@ export function SettingsModal({ open, onOpenChange }: SettingsModalProps) {
 
   const userEmail = session?.user?.email
   const userId = session?.user?.id
+
   const userRole = getUserRole(activeOrganization, userEmail)
   const isOwner = userRole === 'owner'
   const isAdmin = userRole === 'admin'
diff --git a/apps/sim/lib/auth/auth.ts b/apps/sim/lib/auth/auth.ts
@@ -2184,8 +2184,22 @@ export const auth = betterAuth({
                   status: subscription.status,
                 })
 
-                const resolvedSubscription =
-                  await ensureOrganizationForTeamSubscription(subscription)
+                let resolvedSubscription = subscription
+                try {
+                  resolvedSubscription = await ensureOrganizationForTeamSubscription(subscription)
+                } catch (orgError) {
+                  logger.error(
+                    '[onSubscriptionComplete] Failed to ensure organization for team subscription',
+                    {
+                      subscriptionId: subscription.id,
+                      referenceId: subscription.referenceId,
+                      plan: subscription.plan,
+                      error: orgError instanceof Error ? orgError.message : String(orgError),
+                      stack: orgError instanceof Error ? orgError.stack : undefined,
+                    }
+                  )
+                  throw orgError
+                }
 
                 await handleSubscriptionCreated(resolvedSubscription)
 
@@ -2206,8 +2220,22 @@ export const auth = betterAuth({
                   plan: subscription.plan,
                 })
 
-                const resolvedSubscription =
-                  await ensureOrganizationForTeamSubscription(subscription)
+                let resolvedSubscription = subscription
+                try {
+                  resolvedSubscription = await ensureOrganizationForTeamSubscription(subscription)
+                } catch (orgError) {
+                  logger.error(
+                    '[onSubscriptionUpdate] Failed to ensure organization for team subscription',
+                    {
+                      subscriptionId: subscription.id,
+                      referenceId: subscription.referenceId,
+                      plan: subscription.plan,
+                      error: orgError instanceof Error ? orgError.message : String(orgError),
+                      stack: orgError instanceof Error ? orgError.stack : undefined,
+                    }
+                  )
+                  throw orgError
+                }
 
                 try {
                   await syncSubscriptionUsageLimits(resolvedSubscription)
diff --git a/apps/sim/lib/billing/client/upgrade.ts b/apps/sim/lib/billing/client/upgrade.ts
@@ -81,12 +81,15 @@ export function useSubscriptionUpgrade() {
       }
 
       const currentUrl = `${window.location.origin}${window.location.pathname}`
+      const successUrlObj = new URL(window.location.href)
+      successUrlObj.searchParams.set('upgraded', 'true')
+      const successUrl = successUrlObj.toString()
 
       try {
         const upgradeParams = {
           plan: targetPlan,
           referenceId,
-          successUrl: currentUrl,
+          successUrl,
           cancelUrl: currentUrl,
           ...(targetPlan === 'team' && { seats: CONSTANTS.INITIAL_TEAM_SEATS }),
         } as const
diff --git a/apps/sim/lib/billing/organization.ts b/apps/sim/lib/billing/organization.ts
