fix(billing): scope subscription limit syncs to the exact payer (#7695)

* fix(billing): scope subscription limit syncs to the exact payer

* fix(billing): retry subscription limit reconciliation through webhook events

* chore(tests): remove timing-dependent search setup case

Author: icecrasher321

apps/sim/app/workspace/[workspaceId]/search/components/search-source-setup.test.tsx

@@ -413,47 +413,6 @@ describe('organization setup entry points', () => {
     expect(mocks.push).toHaveBeenCalledWith('/o/org-1/settings/integrations/sources/new-source')
   })
 
-  it('honors explicit member-source URLs and clears both setup parameters on close', async () => {
-    useConnectorSetupStore
-      .getState()
-      .saveDraft('user-1:organization:org-1:kb-search:github:members', {
-        sourceConfig: { repository: 'acme/docs' },
-        canonicalModes: {},
-        accessMode: 'members',
-        credentialId: 'cred-source',
-        contentCredentialId: null,
-        disabledTagIds: [],
-        savedAt: Date.now(),
-      })
-    await render(organizationSetup(), '?addConnector=github&source-access=members&search=keep')
-
-    expect(mocks.replace).not.toHaveBeenCalled()
-    expect(document.querySelector('button[aria-label="Choose another source"]')).toBeNull()
-    expect(document.body.textContent).not.toContain('Sync using')
-    expect(document.body.textContent).toContain('Sync documents with')
-    expect(button('Add source')).toBeEnabled()
-    await click(button('Add source'))
-    expect(mocks.create).toHaveBeenCalledWith(
-      expect.objectContaining({
-        connectorType: 'github',
-        accessMode: 'members',
-        sourceConfig: { repository: 'acme/docs' },
-      }),
-      expect.any(Object)
-    )
-    await click(button('Cancel'))
-    expect(mocks.urlUpdate).toHaveBeenLastCalledWith(
-      expect.objectContaining({ queryString: '?search=keep' })
-    )
-    expect(document.querySelector('[role="dialog"]')).toBeNull()
-    expect(mocks.push).not.toHaveBeenCalled()
-    expect(
-      useConnectorSetupStore
-        .getState()
-        .getDraft('user-1:organization:org-1:kb-search:github:members')
-    ).toBeUndefined()
-  })
-
   it.each(['github', 'gmail', 'google_calendar', 'jira'])(
     'returns old %s organization setup links to personal integrations without loading the index',
     async (type) => {

apps/sim/lib/auth/auth.ts

@@ -101,6 +101,7 @@ import {
   handleSubscriptionCreated,
   handleSubscriptionDeleted,
 } from '@/lib/billing/webhooks/subscription'
+import { handleSubscriptionUsageUpdate } from '@/lib/billing/webhooks/subscription-usage'
 import { env } from '@/lib/core/config/env'
 import {
   isAuthDisabled,
@@ -1615,16 +1616,6 @@ export const auth = betterAuth({
                   throw orgError
                 }
 
-                try {
-                  await syncSubscriptionUsageLimits(resolvedSubscription)
-                } catch (error) {
-                  logger.error('[onSubscriptionUpdate] Failed to sync usage limits', {
-                    subscriptionId: resolvedSubscription.id,
-                    referenceId: resolvedSubscription.referenceId,
-                    error,
-                  })
-                }
-
                 if (isTeam(effectivePlanForTeamFeatures)) {
                   try {
                     const quantity = stripeSubscription.items?.data?.[0]?.quantity || 1
@@ -1703,6 +1694,7 @@ export const auth = betterAuth({
                   case 'customer.subscription.created':
                   case 'customer.subscription.updated': {
                     await handleManualEnterpriseSubscription(event)
+                    await handleSubscriptionUsageUpdate(event)
                     break
                   }
                   case 'checkout.session.expired': {

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

@@ -1,8 +1,15 @@
 /**
  * @vitest-environment node
  */
-import { dbChainMockFns, resetEnvFlagsMock, setEnvFlags } from '@sim/testing'
-import { afterAll, beforeAll, beforeEach, describe, expect, it, vi } from 'vitest'
+import {
+  dbChainMockFns,
+  queueTableRows,
+  resetDbChainMock,
+  resetEnvFlagsMock,
+  schemaMock,
+  setEnvFlags,
+} from '@sim/testing'
+import { afterAll, afterEach, beforeAll, beforeEach, describe, expect, it, vi } from 'vitest'
 
 const {
   mockGetHighestPrioritySubscription,
@@ -190,20 +197,39 @@ describe('getOrganizationCoverageForMember', () => {
 describe('getOrganizationIdForSubscriptionReference', () => {
   beforeEach(() => {
     vi.clearAllMocks()
+    resetDbChainMock()
   })
 
-  it('returns an organization id directly when the reference already points to one', async () => {
-    dbChainMockFns.limit.mockResolvedValueOnce([{ id: 'org-1' }])
+  afterEach(resetDbChainMock)
 
-    await expect(getOrganizationIdForSubscriptionReference('org-1')).resolves.toBe('org-1')
-  })
+  it.each(['org-1', 'legacy-organization-id'])(
+    'returns the directly referenced organization %s',
+    async (organizationId) => {
+      queueTableRows(schemaMock.organization, [{ id: organizationId }])
+
+      await expect(getOrganizationIdForSubscriptionReference(organizationId)).resolves.toBe(
+        organizationId
+      )
+    }
+  )
+
+  it.each(['owner', 'admin', 'member'])(
+    'keeps a personal subscription personal when its user is an organization %s',
+    async (role) => {
+      queueTableRows(schemaMock.organization, [])
+      queueTableRows(schemaMock.member, [{ organizationId: 'org-1', role }])
 
-  it('falls back to the admin-owned organization when the reference is still user-scoped', async () => {
-    dbChainMockFns.limit
-      .mockResolvedValueOnce([])
-      .mockResolvedValueOnce([{ organizationId: 'org-1', role: 'owner' }])
+      await expect(getOrganizationIdForSubscriptionReference('user-1')).resolves.toBeNull()
+      expect(dbChainMockFns.from).not.toHaveBeenCalledWith(schemaMock.member)
+    }
+  )
 
-    await expect(getOrganizationIdForSubscriptionReference('user-1')).resolves.toBe('org-1')
+  it('propagates lookup errors instead of treating the subscription as personal', async () => {
+    dbChainMockFns.limit.mockRejectedValueOnce(new Error('db unavailable'))
+
+    await expect(getOrganizationIdForSubscriptionReference('org-1')).rejects.toThrow(
+      'db unavailable'
+    )
   })
 })
 

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

@@ -2,7 +2,6 @@ import { cache } from 'react'
 import { db } from '@sim/db'
 import { member, organization, subscription, user } from '@sim/db/schema'
 import { createLogger } from '@sim/logger'
-import { isOrgAdminRole } from '@sim/platform-authz/workspace'
 import { and, eq, inArray, sql } from 'drizzle-orm'
 import { getEffectiveBillingStatus, isOrganizationBillingBlocked } from '@/lib/billing/core/access'
 import {
@@ -275,6 +274,7 @@ export async function getOrganizationCoverageForMember(
   }
 }
 
+/** Resolves the subscription's exact organization reference without inferring ownership from membership. */
 export async function getOrganizationIdForSubscriptionReference(
   referenceId: string
 ): Promise<string | null> {
@@ -284,24 +284,7 @@ export async function getOrganizationIdForSubscriptionReference(
     .where(eq(organization.id, referenceId))
     .limit(1)
 
-  if (referencedOrganization) {
-    return referencedOrganization.id
-  }
-
-  const [memberRecord] = await db
-    .select({
-      organizationId: member.organizationId,
-      role: member.role,
-    })
-    .from(member)
-    .where(eq(member.userId, referenceId))
-    .limit(1)
-
-  if (memberRecord && isOrgAdminRole(memberRecord.role)) {
-    return memberRecord.organizationId
-  }
-
-  return null
+  return referencedOrganization?.id ?? null
 }
 
 /**

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

@@ -1,10 +1,9 @@
 /**
  * Tests for getUserUsageLimit.
  *
- * Org-scoped members carry a null `currentUsageLimit` by design, so a user
- * whose subscription stops being org-scoped without a resync is left null.
- * The limit read must self-heal that state to the plan/free base plus prepaid
- * balance instead of failing closed and blocking every execution.
+ * Legacy membership syncs may leave a null personal usage limit. The limit
+ * read must recover the plan/free base plus prepaid balance, and subsequent
+ * subscription syncs must preserve independent personal and organization pools.
  *
  * @vitest-environment node
  */
@@ -25,12 +24,14 @@ afterAll(() => {
 const {
   mockGetFreeTierLimit,
   mockGetHighestPrioritySubscription,
+  mockGetHighestPriorityPersonalSubscription,
   mockGetPerUserMinimumLimit,
   mockHasPaidSubscriptionStatus,
   mockIsOrgScopedSubscription,
 } = vi.hoisted(() => ({
   mockGetFreeTierLimit: vi.fn(),
   mockGetHighestPrioritySubscription: vi.fn(),
+  mockGetHighestPriorityPersonalSubscription: vi.fn(),
   mockGetPerUserMinimumLimit: vi.fn(),
   mockHasPaidSubscriptionStatus: vi.fn(),
   mockIsOrgScopedSubscription: vi.fn(),
@@ -48,6 +49,7 @@ vi.mock('@/lib/billing/subscriptions/utils', () => ({
 
 vi.mock('@/lib/billing/core/plan', () => ({
   getHighestPrioritySubscription: mockGetHighestPrioritySubscription,
+  getHighestPriorityPersonalSubscription: mockGetHighestPriorityPersonalSubscription,
 }))
 
 vi.mock('@/lib/billing/core/access', () => ({
@@ -205,10 +207,49 @@ describe('syncUsageLimitsFromSubscription', () => {
     vi.clearAllMocks()
     resetDbChainMock()
     mockIsOrgScopedSubscription.mockReturnValue(false)
+    mockHasPaidSubscriptionStatus.mockImplementation((status: string) => status === 'active')
+  })
+
+  it.each([
+    { plan: 'pro', minimum: 40 },
+    { plan: 'enterprise', minimum: 0 },
+  ])(
+    'preserves a personal $plan cap when the user also belongs to an enterprise organization',
+    async ({ plan, minimum }) => {
+      const personalSubscription = { plan, referenceId: 'user-1', status: 'active' }
+      mockGetHighestPriorityPersonalSubscription.mockResolvedValue(personalSubscription)
+      mockGetHighestPrioritySubscription.mockResolvedValue({
+        plan: 'enterprise',
+        referenceId: 'org-1',
+        status: 'active',
+      })
+      mockIsOrgScopedSubscription.mockReturnValue(true)
+      mockGetPerUserMinimumLimit.mockReturnValue(minimum)
+      dbChainMockFns.limit.mockResolvedValueOnce([{ currentUsageLimit: '80', creditBalance: '1' }])
+
+      await syncUsageLimitsFromSubscription('user-1')
+
+      expect(mockGetHighestPriorityPersonalSubscription).toHaveBeenCalledExactlyOnceWith('user-1', {
+        onError: 'throw',
+      })
+      expect(mockGetHighestPrioritySubscription).not.toHaveBeenCalled()
+      expect(mockGetPerUserMinimumLimit).toHaveBeenCalledWith(personalSubscription)
+      const update = dbChainMockFns.set.mock.calls[0]?.[0]
+      expect(update?.currentUsageLimit).not.toBeNull()
+      expect(JSON.stringify(update?.currentUsageLimit)).toContain('greatest')
+    }
+  )
+
+  it('does not reset a personal cap when its subscription lookup fails', async () => {
+    mockGetHighestPriorityPersonalSubscription.mockRejectedValueOnce(new Error('db unavailable'))
+    dbChainMockFns.limit.mockResolvedValueOnce([{ currentUsageLimit: '80' }])
+
+    await expect(syncUsageLimitsFromSubscription('user-1')).rejects.toThrow('db unavailable')
+    expect(dbChainMockFns.update).not.toHaveBeenCalled()
   })
 
   it('raises a paid personal limit to plan base plus the exact prepaid balance', async () => {
-    mockGetHighestPrioritySubscription.mockResolvedValue(PRO_SUBSCRIPTION)
+    mockGetHighestPriorityPersonalSubscription.mockResolvedValue(PRO_SUBSCRIPTION)
     mockGetPerUserMinimumLimit.mockReturnValue(40)
     dbChainMockFns.limit.mockResolvedValueOnce([
       { currentUsageLimit: '40', creditBalance: '0.005' },
@@ -224,7 +265,7 @@ describe('syncUsageLimitsFromSubscription', () => {
   })
 
   it('restores free-tier base plus prepaid after a downgrade or org departure', async () => {
-    mockGetHighestPrioritySubscription.mockResolvedValue(null)
+    mockGetHighestPriorityPersonalSubscription.mockResolvedValue(null)
     mockGetPerUserMinimumLimit.mockReturnValue(10)
     dbChainMockFns.limit.mockResolvedValueOnce([
       { currentUsageLimit: null, creditBalance: '0.006' },
@@ -240,7 +281,7 @@ describe('syncUsageLimitsFromSubscription', () => {
   })
 
   it('does not retain a higher paid custom cap after downgrade to free', async () => {
-    mockGetHighestPrioritySubscription.mockResolvedValue(null)
+    mockGetHighestPriorityPersonalSubscription.mockResolvedValue(null)
     mockGetPerUserMinimumLimit.mockReturnValue(10)
     dbChainMockFns.limit.mockResolvedValueOnce([
       { currentUsageLimit: '100', creditBalance: '0.006' },
@@ -256,7 +297,7 @@ describe('syncUsageLimitsFromSubscription', () => {
   })
 
   it('preserves a higher custom personal limit', async () => {
-    mockGetHighestPrioritySubscription.mockResolvedValue(PRO_SUBSCRIPTION)
+    mockGetHighestPriorityPersonalSubscription.mockResolvedValue(PRO_SUBSCRIPTION)
     mockGetPerUserMinimumLimit.mockReturnValue(40)
     dbChainMockFns.limit.mockResolvedValueOnce([{ currentUsageLimit: '50', creditBalance: '1' }])
 

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

@@ -7,6 +7,7 @@ import { and, eq, isNull, sql } from 'drizzle-orm'
 import { getEffectiveBillingStatus } from '@/lib/billing/core/access'
 import { defaultBillingPeriod } from '@/lib/billing/core/billing-period'
 import {
+  getHighestPriorityPersonalSubscription,
   getHighestPrioritySubscription,
   type HighestPrioritySubscription,
 } from '@/lib/billing/core/plan'
@@ -449,13 +450,11 @@ export async function updateUserUsageLimit(
  * checks). Org-scoped subs return the organization limit;
  * personally-scoped subs return the individual user limit from userStats.
  *
- * Org-scoped members carry a null `currentUsageLimit` by design (see
- * `syncUsageLimitsFromSubscription`). A user whose subscription stops being
- * org-scoped without a resync would otherwise stay null and fail closed on
- * every execution, so a null limit self-heals to the plan/free base plus the
- * exact prepaid balance here. The write-back is best-effort: a limit written
- * concurrently wins, and a failed write still resolves to the fallback
- * instead of blocking execution.
+ * Legacy organization membership syncs may have cleared the personal limit.
+ * A null limit self-heals to the personal plan/free base plus the exact prepaid
+ * balance here. The write-back is best-effort: a limit written concurrently
+ * wins, and a failed write still resolves to the fallback instead of blocking
+ * execution.
  */
 export async function getUserUsageLimit(
   userId: string,
@@ -576,37 +575,19 @@ export async function checkUsageStatus(userId: string): Promise<{
 }
 
 /**
- * Sync usage limits based on subscription changes
+ * Syncs the user's personal billing pool from their exact personal subscription.
+ * Organization subscriptions have a separate pool and never clear personal limits.
  */
 export async function syncUsageLimitsFromSubscription(userId: string): Promise<void> {
   const [subscription, currentUserStats] = await Promise.all([
-    getHighestPrioritySubscription(userId),
+    getHighestPriorityPersonalSubscription(userId, { onError: 'throw' }),
     db.select(userStatsColumns).from(userStats).where(eq(userStats.userId, userId)).limit(1),
   ])
 
   if (currentUserStats.length === 0) {
     throw new Error(`User stats not found for userId: ${userId}`)
   }
 
-  const currentStats = currentUserStats[0]
-
-  if (isOrgScopedSubscription(subscription, userId)) {
-    if (currentStats.currentUsageLimit !== null) {
-      await db
-        .update(userStats)
-        .set({
-          currentUsageLimit: null,
-          usageLimitUpdatedAt: new Date(),
-        })
-        .where(eq(userStats.userId, userId))
-
-      logger.info('Cleared individual limit for org-scoped member', {
-        userId,
-        plan: subscription?.plan,
-      })
-    }
-    return
-  }
   const baseLimit = toDecimal(getPerUserMinimumLimit(subscription)).toString()
   const hasEntitledPersonalSubscription =
     subscription !== null && hasPaidSubscriptionStatus(subscription.status)
@@ -634,7 +615,6 @@ export async function syncUsageLimitsFromSubscription(userId: string): Promise<v
       : 'Reset limit to free-plus-prepaid minimum',
     { userId, baseLimit: Number(baseLimit) }
   )
-  // Keep higher custom limits unchanged only while personal billing is entitled.
 }
 
 /**