Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -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'])(
Comment thread
icecrasher321 marked this conversation as resolved.
'returns old %s organization setup links to personal integrations without loading the index',
async (type) => {
Expand Down
12 changes: 2 additions & 10 deletions apps/sim/lib/auth/auth.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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': {
Expand Down
48 changes: 37 additions & 11 deletions apps/sim/lib/billing/core/subscription.test.ts
Original file line number Diff line number Diff line change
@@ -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,
Expand Down Expand Up @@ -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'
)
})
})

Expand Down
21 changes: 2 additions & 19 deletions apps/sim/lib/billing/core/subscription.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down Expand Up @@ -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> {
Expand All @@ -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
}

/**
Expand Down
57 changes: 49 additions & 8 deletions apps/sim/lib/billing/core/usage.test.ts
Original file line number Diff line number Diff line change
@@ -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
*/
Expand All @@ -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(),
Expand All @@ -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', () => ({
Expand Down Expand Up @@ -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' },
Expand All @@ -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' },
Expand All @@ -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' },
Expand All @@ -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' }])

Expand Down
38 changes: 9 additions & 29 deletions apps/sim/lib/billing/core/usage.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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'
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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([
Comment thread
icecrasher321 marked this conversation as resolved.
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)
Expand Down Expand Up @@ -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.
}

/**
Expand Down
Loading
Loading