chore: import upstream snapshot with attribution

apps/sim/lib/billing/webhooks/subscription.ts

@@ -0,0 +1,560 @@
+import { AuditAction, AuditResourceType, recordAudit } from '@sim/audit'
+import { db } from '@sim/db'
+import { member, subscription } from '@sim/db/schema'
+import { createLogger } from '@sim/logger'
+import { and, eq, inArray, ne } from 'drizzle-orm'
+import { calculateSubscriptionOverage, isSubscriptionOrgScoped } from '@/lib/billing/core/billing'
+import { syncUsageLimitsFromSubscription } from '@/lib/billing/core/usage'
+import { restoreUserProSubscription } from '@/lib/billing/organizations/membership'
+import { isEnterprise, isPaid, isPro, isTeam } from '@/lib/billing/plan-helpers'
+import { requireStripeClient } from '@/lib/billing/stripe-client'
+import { ENTITLED_SUBSCRIPTION_STATUSES } from '@/lib/billing/subscriptions/utils'
+import { stripeWebhookIdempotency } from '@/lib/billing/webhooks/idempotency'
+import {
+  getBilledOverageForSubscription,
+  resetUsageForSubscription,
+} from '@/lib/billing/webhooks/invoices'
+import { captureServerEvent } from '@/lib/posthog/server'
+import { detachOrganizationWorkspaces } from '@/lib/workspaces/organization-workspaces'
+
+const logger = createLogger('StripeSubscriptionWebhooks')
+
+/**
+ * Resolve a real `user.id` to use as the audit actor for a subscription
+ * event. Org-scoped subscriptions resolve to the org owner; personally
+ * scoped subscriptions already reference a user.
+ */
+async function resolveSubscriptionActorId(referenceId: string): Promise<string> {
+  try {
+    const rows = await db
+      .select({ userId: member.userId })
+      .from(member)
+      .where(and(eq(member.organizationId, referenceId), eq(member.role, 'owner')))
+      .limit(1)
+    return rows[0]?.userId ?? referenceId
+  } catch (error) {
+    logger.warn('Failed to resolve subscription actor; falling back to reference id', {
+      referenceId,
+      error,
+    })
+    return referenceId
+  }
+}
+
+/**
+ * Restore personal Pro subscriptions for every member of an organization
+ * when the team/enterprise subscription ends. Errors propagate so the
+ * enclosing webhook handler fails and Stripe retries the delivery.
+ *
+ * `restoreUserProSubscription` is idempotent: already-restored members
+ * are no-ops on retry, so a partial first attempt is safe to re-run.
+ */
+async function restoreMemberProSubscriptions(organizationId: string): Promise<number> {
+  const members = await db
+    .select({ userId: member.userId })
+    .from(member)
+    .where(eq(member.organizationId, organizationId))
+
+  let restoredCount = 0
+
+  for (const m of members) {
+    const result = await restoreUserProSubscription(m.userId)
+    if (result.restored) {
+      restoredCount++
+    }
+  }
+
+  if (restoredCount > 0) {
+    logger.info('Restored Pro subscriptions for team members', {
+      organizationId,
+      restoredCount,
+      totalMembers: members.length,
+    })
+  }
+
+  return restoredCount
+}
+
+export interface OrganizationDormantTransitionResult {
+  restoredProCount: number
+  membersSynced: number
+  workspacesDetached: number
+  organizationRetainsTeamOrEnterprise: boolean
+}
+
+/**
+ * Returns true when the organization is still covered by an **active
+ * Team or Enterprise** subscription other than `excludeSubscriptionId`.
+ * The org keeps its team-owned workspaces only while such a sub exists;
+ * a Pro sub on the org does not count.
+ */
+async function hasOtherActiveTeamOrEnterpriseSubscription(
+  organizationId: string,
+  excludeSubscriptionId: string | null
+): Promise<boolean> {
+  const filters = [
+    eq(subscription.referenceId, organizationId),
+    inArray(subscription.status, ENTITLED_SUBSCRIPTION_STATUSES),
+  ]
+  if (excludeSubscriptionId) {
+    filters.push(ne(subscription.id, excludeSubscriptionId))
+  }
+
+  const rows = await db
+    .select({ plan: subscription.plan })
+    .from(subscription)
+    .where(and(...filters))
+
+  return rows.some((row) => isTeam(row.plan) || isEnterprise(row.plan))
+}
+
+async function transitionOrganizationToDormantState(
+  organizationId: string,
+  triggeringSubscriptionId: string | null
+): Promise<OrganizationDormantTransitionResult> {
+  const memberUserIds = await db
+    .select({ userId: member.userId })
+    .from(member)
+    .where(eq(member.organizationId, organizationId))
+
+  if (await hasOtherActiveTeamOrEnterpriseSubscription(organizationId, triggeringSubscriptionId)) {
+    logger.info(
+      'Skipping dormant transition - another Team/Enterprise subscription still covers this organization',
+      { organizationId, triggeringSubscriptionId }
+    )
+
+    for (const m of memberUserIds) {
+      await syncUsageLimitsFromSubscription(m.userId)
+    }
+
+    return {
+      restoredProCount: 0,
+      membersSynced: memberUserIds.length,
+      workspacesDetached: 0,
+      organizationRetainsTeamOrEnterprise: true,
+    }
+  }
+
+  const { detachedWorkspaceIds } = await detachOrganizationWorkspaces(organizationId)
+  const restoredProCount = await restoreMemberProSubscriptions(organizationId)
+
+  for (const m of memberUserIds) {
+    await syncUsageLimitsFromSubscription(m.userId)
+  }
+
+  return {
+    restoredProCount,
+    membersSynced: memberUserIds.length,
+    workspacesDetached: detachedWorkspaceIds.length,
+    organizationRetainsTeamOrEnterprise: false,
+  }
+}
+
+export async function handleOrganizationPlanDowngrade(
+  subscriptionData: {
+    id: string
+    plan: string | null
+    referenceId: string
+    status: string | null
+  },
+  stripeEventId?: string
+): Promise<void> {
+  if (!(await isSubscriptionOrgScoped({ referenceId: subscriptionData.referenceId }))) {
+    return
+  }
+
+  const stillTeamOrEnterprise = isTeam(subscriptionData.plan) || isEnterprise(subscriptionData.plan)
+  if (stillTeamOrEnterprise) {
+    return
+  }
+
+  const [currentRow] = await db
+    .select({ plan: subscription.plan })
+    .from(subscription)
+    .where(eq(subscription.id, subscriptionData.id))
+    .limit(1)
+
+  if (currentRow && (isTeam(currentRow.plan) || isEnterprise(currentRow.plan))) {
+    logger.info('Skipping plan downgrade transition - subscription is currently Team/Enterprise', {
+      subscriptionId: subscriptionData.id,
+      organizationId: subscriptionData.referenceId,
+      eventPlan: subscriptionData.plan,
+      currentPlan: currentRow.plan,
+    })
+    return
+  }
+
+  const idempotencyIdentifier = stripeEventId ?? `plan-downgrade:${subscriptionData.id}`
+
+  try {
+    await stripeWebhookIdempotency.executeWithIdempotency(
+      'organization-plan-downgrade',
+      idempotencyIdentifier,
+      async () => {
+        const result = await transitionOrganizationToDormantState(
+          subscriptionData.referenceId,
+          subscriptionData.id
+        )
+
+        if (result.workspacesDetached > 0 || result.restoredProCount > 0) {
+          logger.info('Transitioned organization to dormant state after plan downgrade', {
+            organizationId: subscriptionData.referenceId,
+            subscriptionId: subscriptionData.id,
+            plan: subscriptionData.plan,
+            ...result,
+          })
+        }
+
+        return result
+      }
+    )
+  } catch (error) {
+    logger.error('Failed to transition organization to dormant state on plan downgrade', {
+      organizationId: subscriptionData.referenceId,
+      subscriptionId: subscriptionData.id,
+      plan: subscriptionData.plan,
+      error,
+    })
+    throw error
+  }
+}
+
+/**
+ * Handle new subscription creation - reset usage if transitioning from free to paid
+ */
+export async function handleSubscriptionCreated(
+  subscriptionData: {
+    id: string
+    referenceId: string
+    plan: string | null
+    status: string
+    periodStart?: Date | null
+    periodEnd?: Date | null
+  },
+  stripeEventId?: string
+) {
+  const idempotencyIdentifier = stripeEventId ?? `sub-created:${subscriptionData.id}`
+
+  try {
+    await stripeWebhookIdempotency.executeWithIdempotency(
+      'subscription-created',
+      idempotencyIdentifier,
+      async () => {
+        const otherActiveSubscriptions = await db
+          .select()
+          .from(subscription)
+          .where(
+            and(
+              eq(subscription.referenceId, subscriptionData.referenceId),
+              inArray(subscription.status, ENTITLED_SUBSCRIPTION_STATUSES),
+              ne(subscription.id, subscriptionData.id) // Exclude current subscription
+            )
+          )
+
+        const wasFreePreviously = otherActiveSubscriptions.length === 0
+        const isPaidPlan = isPaid(subscriptionData.plan)
+
+        if (wasFreePreviously && isPaidPlan) {
+          logger.info('Detected free -> paid transition, resetting usage', {
+            subscriptionId: subscriptionData.id,
+            referenceId: subscriptionData.referenceId,
+            plan: subscriptionData.plan,
+          })
+
+          await resetUsageForSubscription({
+            plan: subscriptionData.plan,
+            referenceId: subscriptionData.referenceId,
+            periodStart: subscriptionData.periodStart ?? null,
+            periodEnd: subscriptionData.periodEnd ?? null,
+          })
+
+          logger.info('Successfully reset usage for free -> paid transition', {
+            subscriptionId: subscriptionData.id,
+            referenceId: subscriptionData.referenceId,
+            plan: subscriptionData.plan,
+          })
+        } else {
+          logger.info('No usage reset needed', {
+            subscriptionId: subscriptionData.id,
+            referenceId: subscriptionData.referenceId,
+            plan: subscriptionData.plan,
+            wasFreePreviously,
+            isPaidPlan,
+            otherActiveSubscriptionsCount: otherActiveSubscriptions.length,
+          })
+        }
+
+        if (wasFreePreviously && isPaidPlan) {
+          // Best-effort instrumentation; a transient DB error here must never abort
+          // the (already-committed) free -> paid usage reset above, so it's guarded.
+          try {
+            const actorId = await resolveSubscriptionActorId(subscriptionData.referenceId)
+            const isOrgScoped = await isSubscriptionOrgScoped({
+              referenceId: subscriptionData.referenceId,
+            })
+            recordAudit({
+              actorId,
+              action: AuditAction.SUBSCRIPTION_CREATED,
+              resourceType: AuditResourceType.SUBSCRIPTION,
+              resourceId: subscriptionData.id,
+              description: `Subscription created on ${subscriptionData.plan ?? 'unknown'} plan for ${subscriptionData.referenceId}`,
+              metadata: {
+                plan: subscriptionData.plan,
+                status: subscriptionData.status,
+                referenceId: subscriptionData.referenceId,
+                ...(isOrgScoped ? { organizationId: subscriptionData.referenceId } : {}),
+              },
+            })
+            captureServerEvent(subscriptionData.referenceId, 'subscription_created', {
+              plan: subscriptionData.plan ?? 'unknown',
+              status: subscriptionData.status,
+              reference_id: subscriptionData.referenceId,
+            })
+          } catch (instrumentationError) {
+            logger.warn('Failed to record subscription-created instrumentation', {
+              subscriptionId: subscriptionData.id,
+              referenceId: subscriptionData.referenceId,
+              error: instrumentationError,
+            })
+          }
+        }
+      }
+    )
+  } catch (error) {
+    logger.error('Failed to handle subscription creation usage reset', {
+      subscriptionId: subscriptionData.id,
+      referenceId: subscriptionData.referenceId,
+      error,
+    })
+    throw error
+  }
+}
+
+/**
+ * Handles a subscription deletion (cancel) event. Bills any final-period
+ * overages, resets usage, and transitions the organization to a dormant
+ * state via `transitionOrganizationToDormantState` — the same path used
+ * by plan downgrades. Wrapped in `stripeWebhookIdempotency` so duplicate
+ * event deliveries collapse to one execution; if any step throws, the
+ * webhook retries from scratch.
+ */
+export async function handleSubscriptionDeleted(
+  subscription: {
+    id: string
+    plan: string | null
+    referenceId: string
+    stripeSubscriptionId: string | null
+    seats?: number | null
+    periodStart?: Date | null
+    periodEnd?: Date | null
+  },
+  stripeEventId?: string
+) {
+  const stripeSubscriptionId = subscription.stripeSubscriptionId || ''
+
+  logger.info('Processing subscription deletion', {
+    stripeEventId,
+    stripeSubscriptionId,
+    subscriptionId: subscription.id,
+  })
+
+  // Fall back to the subscription DB id when we don't have an event id
+  // (e.g. called outside the Stripe webhook context). Still dedupes a
+  // single subscription's deletion, just not event-granular.
+  const idempotencyIdentifier = stripeEventId ?? `sub:${subscription.id}`
+
+  try {
+    await stripeWebhookIdempotency.executeWithIdempotency(
+      'subscription-deleted',
+      idempotencyIdentifier,
+      async () => {
+        const totalOverage = await calculateSubscriptionOverage(subscription)
+        const stripe = requireStripeClient()
+
+        if (isEnterprise(subscription.plan)) {
+          await resetUsageForSubscription({
+            plan: subscription.plan,
+            referenceId: subscription.referenceId,
+            periodStart: subscription.periodStart ?? null,
+            periodEnd: subscription.periodEnd ?? null,
+          })
+
+          const dormantResult = await transitionOrganizationToDormantState(
+            subscription.referenceId,
+            subscription.id
+          )
+
+          logger.info('Successfully processed enterprise subscription cancellation', {
+            subscriptionId: subscription.id,
+            stripeSubscriptionId,
+            ...dormantResult,
+          })
+
+          const enterpriseActorId = await resolveSubscriptionActorId(subscription.referenceId)
+          recordAudit({
+            actorId: enterpriseActorId,
+            action: AuditAction.SUBSCRIPTION_CANCELLED,
+            resourceType: AuditResourceType.SUBSCRIPTION,
+            resourceId: subscription.id,
+            description: `Enterprise subscription cancelled for ${subscription.referenceId}`,
+            metadata: {
+              plan: subscription.plan,
+              referenceId: subscription.referenceId,
+              organizationId: subscription.referenceId,
+              kind: 'enterprise',
+            },
+          })
+          captureServerEvent(subscription.referenceId, 'subscription_cancelled', {
+            plan: subscription.plan ?? 'unknown',
+            reference_id: subscription.referenceId,
+          })
+
+          return { totalOverage: 0, kind: 'enterprise' as const }
+        }
+
+        const billedOverage = await getBilledOverageForSubscription(subscription)
+        const remainingOverage = Math.max(0, totalOverage - billedOverage)
+
+        logger.info('Subscription deleted overage calculation', {
+          subscriptionId: subscription.id,
+          totalOverage,
+          billedOverage,
+          remainingOverage,
+        })
+
+        if (remainingOverage > 0 && stripeSubscriptionId) {
+          const stripeSubscription = await stripe.subscriptions.retrieve(stripeSubscriptionId)
+          const customerId = stripeSubscription.customer as string
+          const cents = Math.round(remainingOverage * 100)
+          const endedAt = stripeSubscription.ended_at || Math.floor(Date.now() / 1000)
+          const billingPeriod = new Date(endedAt * 1000).toISOString().slice(0, 7)
+
+          const itemIdemKey = `final-overage-item:${customerId}:${stripeSubscriptionId}:${billingPeriod}`
+          const invoiceIdemKey = `final-overage-invoice:${customerId}:${stripeSubscriptionId}:${billingPeriod}`
+          const finalizeIdemKey = `final-overage-finalize:${customerId}:${stripeSubscriptionId}:${billingPeriod}`
+
+          const overageInvoice = await stripe.invoices.create(
+            {
+              customer: customerId,
+              collection_method: 'charge_automatically',
+              auto_advance: true,
+              description: `Final overage charges for ${subscription.plan} subscription (${billingPeriod})`,
+              metadata: {
+                type: 'final_overage_billing',
+                billingPeriod,
+                subscriptionId: stripeSubscriptionId,
+                cancelledAt: stripeSubscription.canceled_at?.toString() || '',
+              },
+            },
+            { idempotencyKey: invoiceIdemKey }
+          )
+
+          await stripe.invoiceItems.create(
+            {
+              customer: customerId,
+              invoice: overageInvoice.id,
+              amount: cents,
+              currency: 'usd',
+              description: `Usage overage for ${subscription.plan} plan (Final billing period)`,
+              metadata: {
+                type: 'final_usage_overage',
+                usage: remainingOverage.toFixed(2),
+                totalOverage: totalOverage.toFixed(2),
+                billedOverage: billedOverage.toFixed(2),
+                billingPeriod,
+              },
+            },
+            { idempotencyKey: itemIdemKey }
+          )
+
+          if (overageInvoice.id) {
+            await stripe.invoices.finalizeInvoice(
+              overageInvoice.id,
+              {},
+              { idempotencyKey: finalizeIdemKey }
+            )
+          }
+
+          logger.info('Created final overage invoice for cancelled subscription', {
+            subscriptionId: subscription.id,
+            stripeSubscriptionId,
+            invoiceId: overageInvoice.id,
+            totalOverage,
+            billedOverage,
+            remainingOverage,
+            cents,
+            billingPeriod,
+          })
+        } else {
+          logger.info('No overage to bill for cancelled subscription', {
+            subscriptionId: subscription.id,
+            plan: subscription.plan,
+          })
+        }
+
+        await resetUsageForSubscription({
+          plan: subscription.plan,
+          referenceId: subscription.referenceId,
+          periodStart: subscription.periodStart ?? null,
+          periodEnd: subscription.periodEnd ?? null,
+        })
+
+        let restoredProCount = 0
+        let membersSynced = 0
+        let workspacesDetached = 0
+
+        const isOrgScoped = await isSubscriptionOrgScoped(subscription)
+        if (isOrgScoped) {
+          const dormantResult = await transitionOrganizationToDormantState(
+            subscription.referenceId,
+            subscription.id
+          )
+          restoredProCount = dormantResult.restoredProCount
+          membersSynced = dormantResult.membersSynced
+          workspacesDetached = dormantResult.workspacesDetached
+        } else if (isPro(subscription.plan)) {
+          await syncUsageLimitsFromSubscription(subscription.referenceId)
+          membersSynced = 1
+        }
+
+        logger.info('Successfully processed subscription cancellation', {
+          subscriptionId: subscription.id,
+          stripeSubscriptionId,
+          plan: subscription.plan,
+          totalOverage,
+          restoredProCount,
+          membersSynced,
+          workspacesDetached,
+        })
+
+        const cancelActorId = await resolveSubscriptionActorId(subscription.referenceId)
+        recordAudit({
+          actorId: cancelActorId,
+          action: AuditAction.SUBSCRIPTION_CANCELLED,
+          resourceType: AuditResourceType.SUBSCRIPTION,
+          resourceId: subscription.id,
+          description: `Subscription cancelled on ${subscription.plan ?? 'unknown'} plan for ${subscription.referenceId}`,
+          metadata: {
+            plan: subscription.plan,
+            referenceId: subscription.referenceId,
+            totalOverage,
+            ...(isOrgScoped ? { organizationId: subscription.referenceId } : {}),
+          },
+        })
+        captureServerEvent(subscription.referenceId, 'subscription_cancelled', {
+          plan: subscription.plan ?? 'unknown',
+          reference_id: subscription.referenceId,
+        })
+
+        return { totalOverage, remainingOverage, restoredProCount, workspacesDetached }
+      }
+    )
+  } catch (error) {
+    logger.error('Failed to handle subscription deletion', {
+      subscriptionId: subscription.id,
+      stripeSubscriptionId,
+      error,
+    })
+    throw error
+  }
+}