diff --git a/apps/sim/app/workspace/[workspaceId]/upgrade/hooks/use-upgrade-state.test.tsx b/apps/sim/app/workspace/[workspaceId]/upgrade/hooks/use-upgrade-state.test.tsx index 188e3a6523f..f2079037598 100644 --- a/apps/sim/app/workspace/[workspaceId]/upgrade/hooks/use-upgrade-state.test.tsx +++ b/apps/sim/app/workspace/[workspaceId]/upgrade/hooks/use-upgrade-state.test.tsx @@ -120,6 +120,22 @@ describe('useUpgradeState', () => { }) }) + it('shows checkout admission failures through the standard error toast', async () => { + mockHandleUpgrade.mockRejectedValueOnce( + new Error('Your subscription payment is still processing.') + ) + + await act(async () => { + root.render() + }) + + await act(async () => { + await currentState?.doUpgrade('team', 25000) + }) + + expect(mockToastError).toHaveBeenCalledWith('Your subscription payment is still processing.') + }) + it('includes the routed workspace when switching the host billing interval', async () => { await act(async () => { root.render() diff --git a/apps/sim/lib/auth/auth.ts b/apps/sim/lib/auth/auth.ts index 50585c6403b..cdcbcd32903 100644 --- a/apps/sim/lib/auth/auth.ts +++ b/apps/sim/lib/auth/auth.ts @@ -38,6 +38,7 @@ import { } from '@/lib/auth/constants' import { getSessionCookieCacheVersion } from '@/lib/auth/security-policy' import { clampExpiryForSession } from '@/lib/auth/session-policy' +import { getActiveOrganizationId } from '@/lib/auth/session-response' import { guardSubscriptionPlanWrites } from '@/lib/auth/stripe-adapter-guard' import { sendPlanWelcomeEmail } from '@/lib/billing' import { @@ -45,6 +46,12 @@ import { authorizeSubscriptionReference, isPersonalCheckoutRequest, } from '@/lib/billing/authorization' +import { + type CheckoutAdmissionClaim, + claimCheckoutAdmission, + releaseCheckoutAdmission, + resolveCheckoutReferenceId, +} from '@/lib/billing/checkout-admission' import { getOrganizationIdForSubscriptionReference, syncSubscriptionPlan, @@ -1001,20 +1008,46 @@ export const auth = betterAuth({ /** * Personal checkout guard. The Stripe plugin's `authorizeReference` * only runs for organization references (it skips references equal to - * the session user), so duplicate-coverage enforcement for personal - * checkouts lives here: a member of an org with an entitled paid - * subscription must not buy a personal plan on top of it. + * the session user), so personal checkout admission lives here. It + * prevents both a duplicate checkout while Stripe payment is pending + * and a personal plan for someone already covered by an organization. */ if (isBillingEnabled && ctx.path === '/subscription/upgrade') { const session = await getSessionFromCtx(ctx) const sessionUserId = session?.user?.id - if (sessionUserId && isPersonalCheckoutRequest(ctx.body ?? {}, sessionUserId)) { - await assertPersonalCheckoutAllowed(sessionUserId) + if (sessionUserId) { + const requestBody = ctx.body ?? {} + const referenceId = resolveCheckoutReferenceId( + requestBody, + sessionUserId, + getActiveOrganizationId(session) + ) + if (referenceId) { + const checkoutAdmissionClaim = await claimCheckoutAdmission(referenceId) + try { + if (isPersonalCheckoutRequest(requestBody, sessionUserId)) { + await assertPersonalCheckoutAllowed(sessionUserId) + } + } catch (error) { + await releaseCheckoutAdmission(checkoutAdmissionClaim) + throw error + } + return { context: { billingCheckoutAdmissionClaim: checkoutAdmissionClaim } } + } } } return }), + after: createAuthMiddleware(async (ctx) => { + if (!isBillingEnabled || ctx.path !== '/subscription/upgrade') return + const checkoutContext = ctx as typeof ctx & { + billingCheckoutAdmissionClaim?: CheckoutAdmissionClaim + } + if (checkoutContext.billingCheckoutAdmissionClaim) { + await releaseCheckoutAdmission(checkoutContext.billingCheckoutAdmissionClaim) + } + }), }, plugins: [ ...(env.TURNSTILE_SECRET_KEY diff --git a/apps/sim/lib/auth/stripe-adapter-guard.test.ts b/apps/sim/lib/auth/stripe-adapter-guard.test.ts index 824921f0a71..fea713ff40c 100644 --- a/apps/sim/lib/auth/stripe-adapter-guard.test.ts +++ b/apps/sim/lib/auth/stripe-adapter-guard.test.ts @@ -22,7 +22,12 @@ function createBaseAdapter() { const asAdapter = (base: ReturnType) => guardSubscriptionPlanWrites(base as unknown as Parameters[0]) -const ORG_ROW = { id: 'sub-1', referenceId: 'org-1', plan: 'team_6000' } +const ORG_ROW = { + id: 'sub-1', + referenceId: 'org-1', + plan: 'team_6000', + stripeSubscriptionId: 'stripe-sub-1', +} const WHERE = [{ field: 'id', value: 'sub-1' }] describe('guardSubscriptionPlanWrites', () => { @@ -144,6 +149,87 @@ describe('guardSubscriptionPlanWrites', () => { expect(base.update).toHaveBeenCalled() }) + it('blocks rebinding a personal subscription to a different Stripe subscription', async () => { + const base = createBaseAdapter() + const personalPro = { + id: 'personal-pro', + referenceId: 'user-1', + plan: 'pro', + stripeSubscriptionId: 'sub_personal_pro', + } + base.findOne.mockResolvedValueOnce(personalPro) + + const guarded = asAdapter(base) + await expect( + guarded.update({ + model: 'subscription', + where: [{ field: 'id', value: personalPro.id }] as never, + update: { + stripeSubscriptionId: 'sub_enterprise', + status: 'active', + periodEnd: new Date('2026-09-11T18:36:09Z'), + billingInterval: 'month', + }, + }) + ).rejects.toThrow(/already bound to Stripe subscription sub_personal_pro/) + + expect(base.update).not.toHaveBeenCalled() + }) + + it('allows Stripe state updates when the subscription ID is unchanged', async () => { + const base = createBaseAdapter() + base.findOne.mockResolvedValueOnce({ + id: 'personal-pro', + referenceId: 'user-1', + plan: 'pro', + stripeSubscriptionId: 'sub_personal_pro', + }) + + const guarded = asAdapter(base) + await guarded.update({ + model: 'subscription', + where: WHERE as never, + update: { + stripeSubscriptionId: 'sub_personal_pro', + status: 'active', + cancelAtPeriodEnd: true, + }, + }) + + expect(base.update).toHaveBeenCalledWith( + expect.objectContaining({ + update: { + stripeSubscriptionId: 'sub_personal_pro', + status: 'active', + cancelAtPeriodEnd: true, + }, + }) + ) + }) + + it('allows binding an unbound local subscription to Stripe', async () => { + const base = createBaseAdapter() + base.findOne.mockResolvedValueOnce({ + id: 'new-subscription', + referenceId: 'user-1', + plan: 'pro', + stripeSubscriptionId: null, + }) + + const guarded = asAdapter(base) + await guarded.update({ + model: 'subscription', + where: WHERE as never, + update: { stripeSubscriptionId: 'sub_new', status: 'incomplete' }, + }) + + expect(base.update).toHaveBeenCalledWith( + expect.objectContaining({ + update: { stripeSubscriptionId: 'sub_new', status: 'incomplete' }, + }) + ) + }) + it('rejects creating an org-referenced subscription with a non-org plan', async () => { const base = createBaseAdapter() dbChainMockFns.limit.mockResolvedValueOnce([{ id: 'org-1' }]) diff --git a/apps/sim/lib/auth/stripe-adapter-guard.ts b/apps/sim/lib/auth/stripe-adapter-guard.ts index 194af9cd6f3..13f20e4166c 100644 --- a/apps/sim/lib/auth/stripe-adapter-guard.ts +++ b/apps/sim/lib/auth/stripe-adapter-guard.ts @@ -17,11 +17,13 @@ type SubscriptionWriteSurface = Pick< /** * The Better Auth Stripe plugin persists webhook state through the raw * database adapter BEFORE invoking our subscription callbacks — including the - * `plan` column resolved from the Stripe price. That makes the adapter the - * only in-process seam that can enforce the billing invariant that - * organization-referenced subscriptions hold Team/Enterprise plans: by the - * time `syncSubscriptionPlan` runs in a callback, the plugin's write has - * already landed. + * Stripe subscription ID and `plan` column resolved from the Stripe price. + * That makes the adapter the only in-process seam that can enforce two billing + * invariants before the write lands: + * + * - an existing subscription cannot be rebound to a different Stripe + * subscription merely because both subscriptions share a customer; + * - organization-referenced subscriptions hold Team/Enterprise plans. * * Checkout admission blocks user-driven violations; this guard blocks the * remaining vector — an operator swapping an org subscription onto a personal @@ -77,11 +79,31 @@ function guardWriteSurface( return adapter.create(data) }, update: async (data) => { - if (data.model === 'subscription' && hasNonOrgPlanWrite(data.update)) { + if (data.model === 'subscription' && needsSubscriptionRowInspection(data.update)) { const row = await adapter.findOne({ model: 'subscription', where: data.where, }) + + if (row && attemptsStripeSubscriptionRebind(row, data.update)) { + const rejectedStripeSubscriptionId = (data.update as { stripeSubscriptionId: string }) + .stripeSubscriptionId + logger.error( + 'Blocked rebinding an existing subscription to a different Stripe subscription', + { + subscriptionId: row.id, + referenceId: row.referenceId, + currentStripeSubscriptionId: row.stripeSubscriptionId, + rejectedStripeSubscriptionId, + } + ) + throw new Error( + `Subscription ${row.id} is already bound to Stripe subscription ${row.stripeSubscriptionId}; refusing to bind ${rejectedStripeSubscriptionId}` + ) + } + + if (!hasNonOrgPlanWrite(data.update)) return adapter.update(data) + const sanitized = await stripPlanWhenOrgReferenced( row ? [row] : [], data.update as Record @@ -110,6 +132,27 @@ interface SubscriptionRowSlice { id: string referenceId: string plan: string + stripeSubscriptionId: string | null +} + +function needsSubscriptionRowInspection(update: unknown): boolean { + return hasStripeSubscriptionIdWrite(update) || hasNonOrgPlanWrite(update) +} + +function hasStripeSubscriptionIdWrite( + update: unknown +): update is { stripeSubscriptionId: unknown } { + return Boolean(update && typeof update === 'object' && 'stripeSubscriptionId' in update) +} + +function attemptsStripeSubscriptionRebind(row: SubscriptionRowSlice, update: unknown): boolean { + if (!hasStripeSubscriptionIdWrite(update)) return false + const incomingStripeSubscriptionId = update.stripeSubscriptionId + return ( + typeof row.stripeSubscriptionId === 'string' && + typeof incomingStripeSubscriptionId === 'string' && + incomingStripeSubscriptionId !== row.stripeSubscriptionId + ) } function hasNonOrgPlanWrite(update: unknown): boolean { diff --git a/apps/sim/lib/billing/authorization.test.ts b/apps/sim/lib/billing/authorization.test.ts index e2843bd47a0..73eb71600c9 100644 --- a/apps/sim/lib/billing/authorization.test.ts +++ b/apps/sim/lib/billing/authorization.test.ts @@ -1,7 +1,7 @@ /** * @vitest-environment node */ -import { resetDbChainMock } from '@sim/testing' +import { dbChainMockFns, hasMockCondition, resetDbChainMock, schemaMock } from '@sim/testing' import { afterAll, beforeEach, describe, expect, it, vi } from 'vitest' const { @@ -96,6 +96,20 @@ describe('authorizeSubscriptionReference', () => { expect(mockIsOwnerOrAdmin).toHaveBeenCalledWith('owner-1', 'org-1') }) + it('blocks an organization checkout while its bound Stripe subscription is incomplete', async () => { + dbChainMockFns.limit.mockResolvedValueOnce([ + { id: 'subscription-1', stripeSubscriptionId: 'sub_pending' }, + ]) + + await expect( + authorizeSubscriptionReference('owner-1', 'org-1', 'upgrade-subscription', 'team_6000') + ).rejects.toThrow(/subscription payment is still processing/) + + expect(mockHasPaidSubscription).not.toHaveBeenCalled() + expect(mockAssertNoUnresolved).not.toHaveBeenCalled() + expect(mockIsOwnerOrAdmin).not.toHaveBeenCalled() + }) + it('rejects an organization checkout for a pro plan — org references only hold Team/Enterprise', async () => { await expect( authorizeSubscriptionReference('owner-1', 'org-1', 'upgrade-subscription', 'pro_6000') @@ -148,6 +162,46 @@ describe('assertPersonalCheckoutAllowed', () => { await expect(assertPersonalCheckoutAllowed('user-1')).resolves.toBeUndefined() }) + it('keeps abandoned, unbound checkout placeholders retryable', async () => { + await assertPersonalCheckoutAllowed('user-1') + + const predicate = dbChainMockFns.where.mock.calls[0]?.[0] + expect( + hasMockCondition( + predicate, + (node) => node.type === 'eq' && node.left === schemaMock.subscription.referenceId + ) + ).toBe(true) + expect( + hasMockCondition( + predicate, + (node) => + node.type === 'eq' && + node.left === schemaMock.subscription.status && + node.right === 'incomplete' + ) + ).toBe(true) + expect( + hasMockCondition( + predicate, + (node) => + node.type === 'isNotNull' && node.column === schemaMock.subscription.stripeSubscriptionId + ) + ).toBe(true) + }) + + it('blocks a personal checkout while its bound Stripe subscription is incomplete', async () => { + dbChainMockFns.limit.mockResolvedValueOnce([ + { id: 'subscription-1', stripeSubscriptionId: 'sub_pending' }, + ]) + + await expect(assertPersonalCheckoutAllowed('user-1')).rejects.toThrow( + /subscription payment is still processing/ + ) + + expect(mockGetOrganizationCoverageForMember).not.toHaveBeenCalled() + }) + it('rejects checkout when an organization subscription already covers the user', async () => { mockGetOrganizationCoverageForMember.mockResolvedValueOnce({ status: 'covered', diff --git a/apps/sim/lib/billing/authorization.ts b/apps/sim/lib/billing/authorization.ts index b6edd341c26..0197b42ecfb 100644 --- a/apps/sim/lib/billing/authorization.ts +++ b/apps/sim/lib/billing/authorization.ts @@ -1,6 +1,8 @@ import { db } from '@sim/db' +import { subscription } from '@sim/db/schema' import { createLogger } from '@sim/logger' import { APIError } from 'better-auth/api' +import { and, eq, isNotNull } from 'drizzle-orm' import { hasPaidSubscription } from '@/lib/billing' import { isOrganizationOwnerOrAdmin } from '@/lib/billing/core/organization' import { getOrganizationCoverageForMember } from '@/lib/billing/core/subscription' @@ -13,6 +15,44 @@ import { isOrgScopedSubscription } from '@/lib/billing/subscriptions/utils' const logger = createLogger('BillingAuthorization') +/** + * Prevent a billing reference from starting a second Stripe Checkout while a + * previously completed Checkout still has a payment in flight. + * + * Better Auth intentionally reuses an unbound `incomplete` placeholder when a + * customer retries an abandoned Checkout. Once Stripe has bound that row to a + * subscription, however, reusing it would create another Stripe subscription + * and make later webhooks race to own the same local row. + */ +async function assertNoBoundIncompleteSubscription(referenceId: string): Promise { + const [pendingSubscription] = await db + .select({ + id: subscription.id, + stripeSubscriptionId: subscription.stripeSubscriptionId, + }) + .from(subscription) + .where( + and( + eq(subscription.referenceId, referenceId), + eq(subscription.status, 'incomplete'), + isNotNull(subscription.stripeSubscriptionId) + ) + ) + .limit(1) + + if (!pendingSubscription) return + + logger.warn('Blocking checkout - Stripe subscription payment is still pending', { + referenceId, + subscriptionId: pendingSubscription.id, + stripeSubscriptionId: pendingSubscription.stripeSubscriptionId, + }) + throw new APIError('CONFLICT', { + message: + 'Your subscription payment is still processing. Wait for it to finish before starting another checkout. If this takes longer than expected, contact support.', + }) +} + /** * Classify a `/subscription/upgrade` request as a personal checkout using * the same reference resolution as the Better Auth Stripe plugin @@ -39,11 +79,9 @@ export function isPersonalCheckoutRequest( /** * Guard for personal (user-referenced) checkouts on `/subscription/upgrade`. * - * A member of an organization with an entitled paid subscription is already - * covered by that org — their usage pools to it and personal Pro - * subscriptions are paused on join — so a personal checkout would bill the - * same human twice. Throws {@link APIError} with a user-facing message; the - * checkout UI surfaces it as-is. + * Blocks both a bound, incomplete Stripe subscription and a member already + * covered by an entitled organization subscription. Throws {@link APIError} + * with a user-facing message; the checkout UI surfaces it as-is. * * Called from the Better Auth `hooks.before` middleware, NOT from * `authorizeReference`: the Stripe plugin skips `authorizeReference` @@ -54,6 +92,8 @@ export function isPersonalCheckoutRequest( * rejected rather than risking a duplicate subscription. */ export async function assertPersonalCheckoutAllowed(userId: string): Promise { + await assertNoBoundIncompleteSubscription(userId) + const coverage = await getOrganizationCoverageForMember(userId) if (coverage.status === 'covered') { @@ -90,6 +130,8 @@ export async function assertPersonalCheckoutAllowed(userId: string): Promise ({ + mockAtomicallyClaim: vi.fn(), + mockRelease: vi.fn(), + mockIdempotencyService: vi.fn(), +})) + +vi.mock('@/lib/core/idempotency/service', () => ({ + IdempotencyService: class MockIdempotencyService { + constructor(options: unknown) { + mockIdempotencyService(options) + } + + atomicallyClaim(...args: unknown[]) { + return mockAtomicallyClaim(...args) + } + + release(...args: unknown[]) { + return mockRelease(...args) + } + }, +})) + +import { + claimCheckoutAdmission, + releaseCheckoutAdmission, + resolveCheckoutReferenceId, +} from '@/lib/billing/checkout-admission' + +describe('checkout admission', () => { + beforeEach(() => { + mockAtomicallyClaim.mockReset() + mockRelease.mockReset() + mockRelease.mockResolvedValue(undefined) + }) + + it('uses durable, short-lived database claims', () => { + expect(mockIdempotencyService).toHaveBeenCalledWith({ + namespace: 'billing-checkout-admission', + ttlSeconds: 120, + inProgressTtlSeconds: 120, + retryFailures: true, + storeResultBody: false, + forceStorage: 'database', + }) + }) + + it('resolves explicit, organization, and personal references like Better Auth', () => { + expect( + resolveCheckoutReferenceId({ referenceId: 'org-explicit' }, 'user-1', 'org-active') + ).toBe('org-explicit') + expect( + resolveCheckoutReferenceId({ customerType: 'organization' }, 'user-1', 'org-active') + ).toBe('org-active') + expect(resolveCheckoutReferenceId({}, 'user-1', 'org-active')).toBe('user-1') + }) + + it('admits only one overlapping checkout for a billing reference', async () => { + mockAtomicallyClaim + .mockResolvedValueOnce({ + claimed: true, + normalizedKey: 'billing-checkout-admission:stripe:org-1', + storageMethod: 'database', + claimToken: 'claim-1', + }) + .mockResolvedValueOnce({ + claimed: false, + normalizedKey: 'billing-checkout-admission:stripe:org-1', + storageMethod: 'database', + existingResult: { status: 'in-progress' }, + }) + + const firstClaim = await claimCheckoutAdmission('org-1') + await expect(claimCheckoutAdmission('org-1')).rejects.toThrow( + /checkout is already being started/ + ) + expect(firstClaim.claimToken).toBe('claim-1') + }) + + it('releases only the claim owned by this request', async () => { + await releaseCheckoutAdmission({ + normalizedKey: 'billing-checkout-admission:stripe:org-1', + storageMethod: 'database', + claimToken: 'claim-1', + }) + + expect(mockRelease).toHaveBeenCalledWith( + 'billing-checkout-admission:stripe:org-1', + 'database', + 'claim-1' + ) + }) + + it('does not replace a successful checkout response when release fails', async () => { + mockRelease.mockRejectedValueOnce(new Error('database unavailable')) + + await expect( + releaseCheckoutAdmission({ + normalizedKey: 'billing-checkout-admission:stripe:org-1', + storageMethod: 'database', + claimToken: 'claim-1', + }) + ).resolves.toBeUndefined() + }) +}) diff --git a/apps/sim/lib/billing/checkout-admission.ts b/apps/sim/lib/billing/checkout-admission.ts new file mode 100644 index 00000000000..dd1b41d43ef --- /dev/null +++ b/apps/sim/lib/billing/checkout-admission.ts @@ -0,0 +1,86 @@ +import { createLogger } from '@sim/logger' +import { getErrorMessage } from '@sim/utils/errors' +import { APIError } from 'better-auth/api' +import { type AtomicClaimResult, IdempotencyService } from '@/lib/core/idempotency/service' + +const logger = createLogger('BillingCheckoutAdmission') + +const CHECKOUT_ADMISSION_LEASE_SECONDS = 2 * 60 + +const checkoutAdmissionLeases = new IdempotencyService({ + namespace: 'billing-checkout-admission', + ttlSeconds: CHECKOUT_ADMISSION_LEASE_SECONDS, + inProgressTtlSeconds: CHECKOUT_ADMISSION_LEASE_SECONDS, + retryFailures: true, + storeResultBody: false, + forceStorage: 'database', +}) + +export interface CheckoutAdmissionClaim { + normalizedKey: string + storageMethod: AtomicClaimResult['storageMethod'] + claimToken: string +} + +/** + * Resolves the billing reference using the same precedence as Better Auth's + * Stripe `referenceMiddleware`: an explicit reference wins, otherwise an + * organization checkout uses the active organization and a personal checkout + * uses the session user. + */ +export function resolveCheckoutReferenceId( + body: { referenceId?: unknown; customerType?: unknown }, + sessionUserId: string, + activeOrganizationId: string | null +): string | null { + if (body.referenceId) { + return typeof body.referenceId === 'string' ? body.referenceId : null + } + if (body.customerType === 'organization') return activeOrganizationId + return sessionUserId +} + +/** + * Atomically reserves one in-flight checkout request per billing reference. + * The claim spans Better Auth's local-row preparation and Stripe Checkout + * creation, closing the window where two requests could both pass a read-only + * admission check before either one bound a subscription. + */ +export async function claimCheckoutAdmission(referenceId: string): Promise { + const claim = await checkoutAdmissionLeases.atomicallyClaim('stripe', referenceId) + if (!claim.claimed) { + logger.warn('Blocking concurrent subscription checkout', { referenceId }) + throw new APIError('CONFLICT', { + message: + 'A subscription checkout is already being started. Please wait a moment and try again.', + }) + } + if (!claim.claimToken) { + throw new Error('Checkout admission claim is missing its fencing token') + } + return { + normalizedKey: claim.normalizedKey, + storageMethod: claim.storageMethod, + claimToken: claim.claimToken, + } +} + +/** + * Releases a checkout admission claim without masking the endpoint result. + * A failed release is bounded by the claim's short lease and remains visible + * in logs for operators. + */ +export async function releaseCheckoutAdmission(claim: CheckoutAdmissionClaim): Promise { + try { + await checkoutAdmissionLeases.release( + claim.normalizedKey, + claim.storageMethod, + claim.claimToken + ) + } catch (error) { + logger.warn('Failed to release checkout admission claim; lease will expire', { + normalizedKey: claim.normalizedKey, + error: getErrorMessage(error, 'Unknown error'), + }) + } +} diff --git a/apps/sim/lib/billing/webhooks/enterprise.test.ts b/apps/sim/lib/billing/webhooks/enterprise.test.ts index 982683f0776..018edda9918 100644 --- a/apps/sim/lib/billing/webhooks/enterprise.test.ts +++ b/apps/sim/lib/billing/webhooks/enterprise.test.ts @@ -15,6 +15,10 @@ const mocks = vi.hoisted(() => ({ subscriptionsRetrieve: vi.fn(), patchOutboxEventPayload: vi.fn(), reapplyPaidOrgJoinBillingForExistingMemberTx: vi.fn(), + acquireUserBillingIdentityLock: vi.fn(), + acquireInvitationMutationLocks: vi.fn(), + attachOwnedWorkspacesToOrganizationTx: vi.fn(), + invalidateWorkspaceTableLimitsCache: vi.fn(), })) vi.mock('@sim/audit', () => ({ @@ -35,6 +39,10 @@ vi.mock('@/lib/billing/organizations/membership', () => ({ reapplyPaidOrgJoinBillingForExistingMemberTx: mocks.reapplyPaidOrgJoinBillingForExistingMemberTx, })) +vi.mock('@/lib/billing/organizations/billing-identity-lock', () => ({ + acquireUserBillingIdentityLock: mocks.acquireUserBillingIdentityLock, +})) + vi.mock('@/lib/billing/stripe-client', () => ({ requireStripeClient: () => ({ subscriptions: { retrieve: mocks.subscriptionsRetrieve }, @@ -64,6 +72,10 @@ vi.mock('@/lib/core/outbox/service', () => ({ patchOutboxEventPayload: mocks.patchOutboxEventPayload, })) +vi.mock('@/lib/invitations/locks', () => ({ + acquireInvitationMutationLocks: mocks.acquireInvitationMutationLocks, +})) + vi.mock('@/lib/messaging/email/mailer', () => ({ sendEmail: vi.fn(), })) @@ -76,6 +88,15 @@ vi.mock('@/lib/posthog/server', () => ({ captureServerEvent: vi.fn(), })) +vi.mock('@/lib/table/billing', () => ({ + invalidateWorkspaceTableLimitsCache: mocks.invalidateWorkspaceTableLimitsCache, +})) + +vi.mock('@/lib/workspaces/organization-workspaces', () => ({ + attachOwnedWorkspacesToOrganizationTx: mocks.attachOwnedWorkspacesToOrganizationTx, + ownedAttachableWorkspacesWhere: vi.fn(), +})) + import { handleManualEnterpriseSubscription } from '@/lib/billing/webhooks/enterprise' const ENTERPRISE_PROVISION_EVENT_TYPE = 'stripe.provision-enterprise' @@ -159,9 +180,18 @@ function eventFor(subscription: Stripe.Subscription): Stripe.Event { function queueSuccessfulExistingSubscriptionReconciliation(options: { operation?: ReturnType + workspaceIds?: string[] }) { queueTableRows(schemaMock.organization, [{ creditBalance: '0' }]) if (options.operation) { + queueTableRows(schemaMock.outboxEvent, [ + { eventType: ENTERPRISE_PROVISION_EVENT_TYPE, payload: options.operation }, + ]) + if (!('applicationResult' in options.operation)) { + const workspaceRows = (options.workspaceIds ?? []).map((id) => ({ id })) + queueTableRows(schemaMock.workspace, workspaceRows) + queueTableRows(schemaMock.workspace, workspaceRows) + } queueTableRows(schemaMock.outboxEvent, [ { eventType: ENTERPRISE_PROVISION_EVENT_TYPE, payload: options.operation }, ]) @@ -180,6 +210,12 @@ describe('Enterprise webhook issuance correlation', () => { resetDbChainMock() mocks.patchOutboxEventPayload.mockResolvedValue(true) mocks.reapplyPaidOrgJoinBillingForExistingMemberTx.mockResolvedValue(undefined) + mocks.attachOwnedWorkspacesToOrganizationTx.mockResolvedValue({ + attachedWorkspaceIds: [], + addedMemberIds: [], + skippedMembers: [], + usageLimitUserIds: [], + }) }) afterAll(() => { @@ -189,6 +225,13 @@ describe('Enterprise webhook issuance correlation', () => { it('retries when the create webhook races ahead of paused-collection provisioning', async () => { const subscription = stripeSubscription({ operationId: 'operation-1', paused: false }) mocks.subscriptionsRetrieve.mockResolvedValue(subscription) + queueTableRows(schemaMock.outboxEvent, [ + { + eventType: ENTERPRISE_PROVISION_EVENT_TYPE, + payload: operationPayload({ pausePaymentCollection: true }), + }, + ]) + queueTableRows(schemaMock.workspace, []) queueTableRows(schemaMock.organization, [{ creditBalance: '0' }]) queueTableRows(schemaMock.outboxEvent, [ { @@ -207,6 +250,66 @@ describe('Enterprise webhook issuance correlation', () => { expect(mocks.reapplyPaidOrgJoinBillingForExistingMemberTx).not.toHaveBeenCalled() }) + it('sweeps the Enterprise owner personal workspaces when issuance is applied', async () => { + const subscription = stripeSubscription({ operationId: 'operation-1', paused: false }) + mocks.subscriptionsRetrieve.mockResolvedValue(subscription) + queueSuccessfulExistingSubscriptionReconciliation({ + operation: operationPayload(), + workspaceIds: ['workspace-1', 'workspace-archived'], + }) + mocks.attachOwnedWorkspacesToOrganizationTx.mockResolvedValueOnce({ + attachedWorkspaceIds: ['workspace-1', 'workspace-archived'], + addedMemberIds: [], + skippedMembers: [], + usageLimitUserIds: [], + }) + + await expect( + handleManualEnterpriseSubscription(eventFor(subscription)) + ).resolves.toBeUndefined() + + expect(mocks.acquireInvitationMutationLocks).toHaveBeenCalledWith(expect.anything(), { + invitationIds: [], + workspaceIds: ['workspace-1', 'workspace-archived'], + }) + expect(mocks.acquireUserBillingIdentityLock).toHaveBeenCalledWith(expect.anything(), 'owner-1') + expect(mocks.attachOwnedWorkspacesToOrganizationTx).toHaveBeenCalledWith(expect.anything(), { + ownerUserId: 'owner-1', + organizationId: 'org-1', + workspaceIds: ['workspace-1', 'workspace-archived'], + externalMemberPolicy: 'external-all', + ownerMatch: 'owner', + includeArchived: true, + }) + expect(mocks.invalidateWorkspaceTableLimitsCache).toHaveBeenCalledTimes(2) + expect(mocks.patchOutboxEventPayload).toHaveBeenCalled() + }) + + it('retries without applying when the Enterprise owner workspace set changes', async () => { + const subscription = stripeSubscription({ operationId: 'operation-1', paused: false }) + mocks.subscriptionsRetrieve.mockResolvedValue(subscription) + queueTableRows(schemaMock.outboxEvent, [ + { eventType: ENTERPRISE_PROVISION_EVENT_TYPE, payload: operationPayload() }, + ]) + queueTableRows(schemaMock.workspace, [{ id: 'workspace-1' }]) + queueTableRows(schemaMock.organization, [{ creditBalance: '0' }]) + queueTableRows(schemaMock.outboxEvent, [ + { eventType: ENTERPRISE_PROVISION_EVENT_TYPE, payload: operationPayload() }, + ]) + queueTableRows(schemaMock.user, [{ stripeCustomerId: 'cus_1' }]) + queueTableRows(schemaMock.member, [{ value: 1 }]) + queueTableRows(schemaMock.subscription, []) + queueTableRows(schemaMock.subscription, [{ id: 'local-sub-1', referenceId: 'org-1' }]) + queueTableRows(schemaMock.workspace, [{ id: 'workspace-1' }, { id: 'workspace-2' }]) + + await expect(handleManualEnterpriseSubscription(eventFor(subscription))).rejects.toThrow( + 'personal workspaces changed during reconciliation' + ) + + expect(mocks.attachOwnedWorkspacesToOrganizationTx).not.toHaveBeenCalled() + expect(mocks.patchOutboxEventPayload).not.toHaveBeenCalled() + }) + it('allows later Stripe metadata edits after the issuance was already applied', async () => { const subscription = stripeSubscription({ operationId: 'operation-1', paused: false }) mocks.subscriptionsRetrieve.mockResolvedValue(subscription) diff --git a/apps/sim/lib/billing/webhooks/enterprise.ts b/apps/sim/lib/billing/webhooks/enterprise.ts index 59c81542064..aa7bc71681d 100644 --- a/apps/sim/lib/billing/webhooks/enterprise.ts +++ b/apps/sim/lib/billing/webhooks/enterprise.ts @@ -1,6 +1,6 @@ import { AuditAction, AuditResourceType, recordAudit } from '@sim/audit' import { db } from '@sim/db' -import { member, organization, outboxEvent, subscription, user } from '@sim/db/schema' +import { member, organization, outboxEvent, subscription, user, workspace } from '@sim/db/schema' import { createLogger } from '@sim/logger' import { generateId } from '@sim/utils/id' import { and, asc, count, eq, inArray, sql } from 'drizzle-orm' @@ -13,6 +13,7 @@ import { enterpriseOperationMatchesStripeSubscription, parseEnterpriseProvisionPayload, } from '@/lib/billing/enterprise-outbox' +import { acquireUserBillingIdentityLock } from '@/lib/billing/organizations/billing-identity-lock' import { acquireOrganizationMutationLock, reapplyPaidOrgJoinBillingForExistingMemberTx, @@ -26,13 +27,67 @@ import { } from '@/lib/billing/webhooks/enterprise-reconciliation-lease' import { stripeWebhookIdempotency } from '@/lib/billing/webhooks/idempotency' import { patchOutboxEventPayload } from '@/lib/core/outbox/service' +import { acquireInvitationMutationLocks } from '@/lib/invitations/locks' import { sendEmail } from '@/lib/messaging/email/mailer' import { getFromEmailAddress } from '@/lib/messaging/email/utils' import { captureServerEvent } from '@/lib/posthog/server' +import { invalidateWorkspaceTableLimitsCache } from '@/lib/table/billing' +import { + attachOwnedWorkspacesToOrganizationTx, + ownedAttachableWorkspacesWhere, +} from '@/lib/workspaces/organization-workspaces' import { parseEnterpriseSubscriptionMetadata } from '../types' const logger = createLogger('BillingEnterprise') +interface EnterpriseWorkspaceSweepPlan { + operationId: string + ownerUserId: string + workspaceIds: string[] +} + +function sameOrderedIds(left: string[], right: string[]): boolean { + return left.length === right.length && left.every((id, index) => id === right[index]) +} + +async function planEnterpriseOwnerWorkspaceSweep( + metadata: Stripe.Metadata, + organizationId: string +): Promise { + const operationId = metadata.enterpriseOperationId + if (!operationId) return null + + const [operationRow] = await db + .select({ eventType: outboxEvent.eventType, payload: outboxEvent.payload }) + .from(outboxEvent) + .where(eq(outboxEvent.id, operationId)) + .limit(1) + const payload = operationRow ? parseEnterpriseProvisionPayload(operationRow.payload) : null + if ( + operationRow?.eventType !== ENTERPRISE_PROVISION_EVENT_TYPE || + !payload || + payload.applicationResult || + payload.request.organizationId !== organizationId + ) { + return null + } + + const workspaceIds = ( + await db + .select({ id: workspace.id }) + .from(workspace) + .where( + ownedAttachableWorkspacesWhere({ + userId: payload.request.ownerUserId, + includeArchived: true, + }) + ) + .orderBy(workspace.id) + ).map((row) => row.id) + + return { operationId, ownerUserId: payload.request.ownerUserId, workspaceIds } +} + export async function handleManualEnterpriseSubscription(event: Stripe.Event) { return stripeWebhookIdempotency.executeWithIdempotency( 'manual-enterprise-subscription', @@ -115,6 +170,7 @@ async function reconcileManualEnterpriseSubscription( } const { seats, monthlyPrice } = enterpriseMetadata + const workspaceSweepPlan = await planEnterpriseOwnerWorkspaceSweep(metadata, referenceId) // Get the first subscription item which contains the period information const referenceItem = stripeSubscription.items?.data?.[0] @@ -148,6 +204,12 @@ async function reconcileManualEnterpriseSubscription( } const coreResult = await db.transaction(async (tx) => { + if (workspaceSweepPlan && workspaceSweepPlan.workspaceIds.length > 0) { + await acquireInvitationMutationLocks(tx, { + invitationIds: [], + workspaceIds: workspaceSweepPlan.workspaceIds, + }) + } await acquireOrganizationMutationLock(tx, referenceId) await tx.execute( sql`select pg_advisory_xact_lock(hashtextextended(${`stripe-subscription:${stripeSubscription.id}`}, 0))` @@ -273,7 +335,7 @@ async function reconcileManualEnterpriseSubscription( if (existing && existing.referenceId !== referenceId) { throw new Error( - `Stripe subscription ${stripeSubscription.id} is already bound to organization ${existing.referenceId}` + `Stripe subscription ${stripeSubscription.id} is already bound to reference ${existing.referenceId}, not organization ${referenceId}` ) } @@ -315,6 +377,46 @@ async function reconcileManualEnterpriseSubscription( }) .where(eq(organization.id, referenceId)) + let attachedWorkspaceIds: string[] = [] + if (operationNewlyApplied && correlatedOperation) { + if ( + !workspaceSweepPlan || + workspaceSweepPlan.operationId !== operationId || + workspaceSweepPlan.ownerUserId !== correlatedOperation.request.ownerUserId + ) { + throw new Error('Unable to establish the Enterprise owner workspace sweep') + } + + await acquireUserBillingIdentityLock(tx, workspaceSweepPlan.ownerUserId) + const currentWorkspaceIds = ( + await tx + .select({ id: workspace.id }) + .from(workspace) + .where( + ownedAttachableWorkspacesWhere({ + userId: workspaceSweepPlan.ownerUserId, + includeArchived: true, + }) + ) + .orderBy(workspace.id) + ).map((row) => row.id) + if (!sameOrderedIds(workspaceSweepPlan.workspaceIds, currentWorkspaceIds)) { + throw new Error( + 'Enterprise owner personal workspaces changed during reconciliation; retry the webhook' + ) + } + + const attached = await attachOwnedWorkspacesToOrganizationTx(tx, { + ownerUserId: workspaceSweepPlan.ownerUserId, + organizationId: referenceId, + workspaceIds: currentWorkspaceIds, + externalMemberPolicy: 'external-all', + ownerMatch: 'owner', + includeArchived: true, + }) + attachedWorkspaceIds = attached.attachedWorkspaceIds + } + // The organization lock is held across the census and all member billing // transitions. Add/remove/accept paths take the same lock, so a departing // member cannot be re-paused after their removal restores personal Pro. @@ -346,6 +448,7 @@ async function reconcileManualEnterpriseSubscription( operationNewlyApplied, hasCorrelatedOperation: Boolean(correlatedOperation), subscriptionNewlyInserted: !existing, + attachedWorkspaceIds, ...creditLimits, } }) @@ -357,10 +460,14 @@ async function reconcileManualEnterpriseSubscription( operationNewlyApplied, hasCorrelatedOperation, subscriptionNewlyInserted, + attachedWorkspaceIds, configuredUsageLimitCredits, prepaidCredits, effectiveUsageLimitCredits, } = coreResult + for (const workspaceId of attachedWorkspaceIds) { + invalidateWorkspaceTableLimitsCache(workspaceId) + } const shouldAnnounce = hasCorrelatedOperation ? operationNewlyApplied : subscriptionNewlyInserted logger.info('[subscription.created] Upserted enterprise subscription', { @@ -372,6 +479,7 @@ async function reconcileManualEnterpriseSubscription( effectiveUsageLimitCredits, prepaidCredits, seats, + attachedWorkspaceCount: attachedWorkspaceIds.length, note: 'Seats from metadata, Stripe quantity set to 1', })