fix(sso): keep subscription admission consistent

Author: waleedlatif1

apps/sim/ee/sso/components/sso-form.tsx

@@ -13,7 +13,7 @@ import { AuthFormMessage, AuthSubmitButton } from '@/app/(auth)/components'
 
 const logger = createLogger('SSOForm')
 const SSO_SIGN_IN_ERROR = 'Unable to start SSO. Check your email and try again.'
-const SSO_ERROR_MESSAGES: Record<string, string> = {
+const SSO_ERROR_MESSAGES = {
   account_not_found: 'No account found. Please contact your administrator to set up SSO access.',
   sso_failed: 'SSO authentication failed. Please try again.',
   invalid_provider: 'SSO provider not configured correctly.',
@@ -23,7 +23,7 @@ const SSO_ERROR_MESSAGES: Record<string, string> = {
     'This Sim account is already a member of another organization. Leave that organization before trying again, or ask an administrator for external workspace access.',
   sso_provisioning_failed:
     'SSO succeeded, but organization access could not be set up safely. No session was created; please try again or contact your administrator.',
-}
+} as const
 
 const validateEmailField = (emailValue: string): string[] => {
   const errors: string[] = []
@@ -50,7 +50,9 @@ export default function SSOForm({ registrationDisabled }: SSOFormProps) {
   const searchParams = useSearchParams()
   const errorCode = searchParams?.get('error') ?? null
   const initialError = errorCode
-    ? SSO_ERROR_MESSAGES[errorCode] || 'SSO authentication failed. Please try again.'
+    ? Object.hasOwn(SSO_ERROR_MESSAGES, errorCode)
+      ? SSO_ERROR_MESSAGES[errorCode as keyof typeof SSO_ERROR_MESSAGES]
+      : 'SSO authentication failed. Please try again.'
     : null
 
   return (

apps/sim/lib/auth/sso/application/admit-sso-user.test.ts

@@ -246,7 +246,7 @@ describe('SSO JIT admission', () => {
 
   it('uses elastic Team seats and reconciles the billed count after admission', async () => {
     queueIdentity()
-    queueTableRows(schemaMock.subscription, [{ plan: 'team' }])
+    queueTableRows(schemaMock.subscription, [{ id: 'subscription-current', plan: 'team' }])
 
     await expect(execute()).resolves.toEqual({
       kind: 'provisioned',
@@ -258,11 +258,30 @@ describe('SSO JIT admission', () => {
       organizationId: 'org-1',
       role: 'member',
       skipSeatValidation: true,
+      organizationSubscriptionId: 'subscription-current',
     })
     expect(mockReconcileOrganizationSeats).toHaveBeenCalledWith({
       organizationId: 'org-1',
       reason: 'sso-jit-member-added',
       actorId: 'user-1',
+      subscriptionId: 'subscription-current',
+    })
+  })
+
+  it('pins fixed Enterprise capacity checks to the subscription used for plan classification', async () => {
+    queueIdentity()
+    queueTableRows(schemaMock.subscription, [{ id: 'subscription-current', plan: 'enterprise' }])
+
+    await expect(execute()).resolves.toEqual({
+      kind: 'provisioned',
+      organizationId: 'org-1',
+      memberId: 'member-1',
+    })
+    expect(mockEnsureUserInOrganizationTx).toHaveBeenCalledWith(dbChainMock.db, {
+      userId: 'user-1',
+      organizationId: 'org-1',
+      role: 'member',
+      organizationSubscriptionId: 'subscription-current',
     })
   })
 

apps/sim/lib/auth/sso/application/admit-sso-user.ts

@@ -64,6 +64,7 @@ interface SuccessfulAdmission {
   result: SsoJitAdmissionResult
   userName?: string
   userEmail?: string
+  organizationSubscriptionId?: string
 }
 
 async function runAdmissionTransaction(
@@ -225,15 +226,15 @@ async function runAdmissionTransaction(
     }
 
     const [organizationSubscription] = await tx
-      .select({ plan: subscription.plan })
+      .select({ id: subscription.id, plan: subscription.plan })
       .from(subscription)
       .where(
         and(
           eq(subscription.referenceId, provider.organizationId),
           inArray(subscription.status, ENTITLED_SUBSCRIPTION_STATUSES)
         )
       )
-      .orderBy(desc(subscription.periodStart))
+      .orderBy(desc(subscription.periodStart), desc(subscription.id))
       .limit(1)
 
     const membershipResult = await ensureUserInOrganizationTx(tx, {
@@ -242,6 +243,9 @@ async function runAdmissionTransaction(
       role: 'member',
       /** Team seats grow to the committed member count; Enterprise remains fixed-capacity. */
       ...(isTeam(organizationSubscription?.plan) ? { skipSeatValidation: true } : {}),
+      ...(organizationSubscription?.id
+        ? { organizationSubscriptionId: organizationSubscription.id }
+        : {}),
     })
 
     if (!membershipResult.success || !membershipResult.memberId) {
@@ -258,6 +262,9 @@ async function runAdmissionTransaction(
 
     return {
       ...attribution,
+      ...(organizationSubscription?.id
+        ? { organizationSubscriptionId: organizationSubscription.id }
+        : {}),
       result: {
         kind: membershipResult.alreadyMember ? 'already-member' : 'provisioned',
         organizationId: provider.organizationId,
@@ -315,6 +322,9 @@ async function runProvisioningPostCommitEffects(
       organizationId,
       reason: 'sso-jit-member-added',
       actorId: userId,
+      ...(admission.organizationSubscriptionId
+        ? { subscriptionId: admission.organizationSubscriptionId }
+        : {}),
     })
   } catch (error) {
     logger.error('Failed to reconcile organization seats after SSO JIT admission', {

apps/sim/lib/billing/organizations/membership.ts

@@ -23,7 +23,7 @@ import { createLogger } from '@sim/logger'
 import { getErrorMessage } from '@sim/utils/errors'
 import { generateId } from '@sim/utils/id'
 import { normalizeEmail } from '@sim/utils/string'
-import { and, count, eq, inArray, isNull, ne, or, sql } from 'drizzle-orm'
+import { and, count, desc, eq, inArray, isNull, ne, or, sql } from 'drizzle-orm'
 import { invalidateMembershipCache } from '@/lib/auth/security-policy'
 import { applySessionPolicyToNewMember } from '@/lib/auth/session-policy'
 import { syncUsageLimitsFromSubscription } from '@/lib/billing/core/usage'
@@ -429,6 +429,8 @@ export interface AddMemberParams {
   skipBillingLogic?: boolean
   /** Skip seat validation (default: false) */
   skipSeatValidation?: boolean
+  /** Restrict billing decisions to an already-resolved entitled organization subscription. */
+  organizationSubscriptionId?: string
   /** When provided, the acceptor's own pending invitation is excluded from the seat count during validation. */
   acceptingInvitationId?: string
 }
@@ -593,6 +595,7 @@ export async function ensureUserInOrganizationTx(
     role,
     skipBillingLogic = false,
     skipSeatValidation = false,
+    organizationSubscriptionId,
   } = params
   const emptyBillingActions = {
     proUsageSnapshotted: false,
@@ -665,9 +668,13 @@ export async function ensureUserInOrganizationTx(
       .where(
         and(
           eq(subscriptionTable.referenceId, organizationId),
-          inArray(subscriptionTable.status, ENTITLED_SUBSCRIPTION_STATUSES)
+          inArray(subscriptionTable.status, ENTITLED_SUBSCRIPTION_STATUSES),
+          organizationSubscriptionId
+            ? eq(subscriptionTable.id, organizationSubscriptionId)
+            : undefined
         )
       )
+      .orderBy(desc(subscriptionTable.periodStart), desc(subscriptionTable.id))
       .limit(1)
     if (!organizationSubscription || !isPaid(organizationSubscription.plan)) {
       return {
@@ -724,9 +731,13 @@ export async function ensureUserInOrganizationTx(
           .where(
             and(
               eq(subscriptionTable.referenceId, organizationId),
-              inArray(subscriptionTable.status, ENTITLED_SUBSCRIPTION_STATUSES)
+              inArray(subscriptionTable.status, ENTITLED_SUBSCRIPTION_STATUSES),
+              organizationSubscriptionId
+                ? eq(subscriptionTable.id, organizationSubscriptionId)
+                : undefined
             )
           )
+          .orderBy(desc(subscriptionTable.periodStart), desc(subscriptionTable.id))
           .limit(1)
         return organizationSubscription && isPaid(organizationSubscription.plan)
           ? applyPaidOrgJoinBillingTx(tx, userId, organizationId)
@@ -2008,6 +2019,7 @@ export async function addUserToOrganization(params: AddMemberParams): Promise<Ad
     role,
     skipBillingLogic = false,
     skipSeatValidation = false,
+    organizationSubscriptionId,
     acceptingInvitationId,
   } = params
 
@@ -2066,6 +2078,7 @@ export async function addUserToOrganization(params: AddMemberParams): Promise<Ad
         role,
         skipBillingLogic,
         skipSeatValidation,
+        organizationSubscriptionId,
         acceptingInvitationId,
       })
     )

apps/sim/lib/billing/organizations/seats.ts

@@ -2,7 +2,7 @@ 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, count, eq, inArray } from 'drizzle-orm'
+import { and, count, desc, eq, inArray } from 'drizzle-orm'
 import { syncSubscriptionUsageLimits } from '@/lib/billing/organization'
 import { isTeam } from '@/lib/billing/plan-helpers'
 import { ENTITLED_SUBSCRIPTION_STATUSES } from '@/lib/billing/subscriptions/utils'
@@ -24,6 +24,8 @@ export interface ReconcileOrganizationSeatsResult {
 interface ReconcileOrganizationSeatsParams {
   organizationId: string
   reason: string
+  /** Restrict reconciliation to an already-resolved entitled organization subscription. */
+  subscriptionId?: string
   /**
    * Real `user.id` of the actor whose action triggered this reconcile, used to
    * attribute the seat-change audit log and analytics event. Omit for system
@@ -50,6 +52,7 @@ interface ReconcileOrganizationSeatsParams {
 export async function reconcileOrganizationSeats({
   organizationId,
   reason,
+  subscriptionId,
   actorId,
 }: ReconcileOrganizationSeatsParams): Promise<ReconcileOrganizationSeatsResult> {
   if (!isBillingEnabled) {
@@ -74,9 +77,11 @@ export async function reconcileOrganizationSeats({
       .where(
         and(
           eq(subscription.referenceId, organizationId),
-          inArray(subscription.status, ENTITLED_SUBSCRIPTION_STATUSES)
+          inArray(subscription.status, ENTITLED_SUBSCRIPTION_STATUSES),
+          subscriptionId ? eq(subscription.id, subscriptionId) : undefined
         )
       )
+      .orderBy(desc(subscription.periodStart), desc(subscription.id))
       .for('update')
       .limit(1)