fix(billing): serialize credit caps with access revocation

Author: waleedlatif1

apps/sim/app/api/organizations/[id]/members/[memberId]/usage-limit/route.test.ts

@@ -2,6 +2,7 @@
  * @vitest-environment node
  */
 
+import { db } from '@sim/db'
 import { member } from '@sim/db/schema'
 import {
   auditMock,
@@ -187,9 +188,10 @@ describe('PUT /api/organizations/[id]/members/[memberId]/usage-limit', () => {
   })
 
   it('persists the limit as dollars (credits / 200) and audits', async () => {
+    queueTableRows(member, [{ role: 'admin' }])
     const res = await PUT(putRequest({ creditLimit: 400 }), context())
     expect(res.status).toBe(200)
-    expect(mockSetOrgMemberUsageLimit).toHaveBeenCalledWith('org-1', 'user-2', 2, 'admin-1')
+    expect(mockSetOrgMemberUsageLimit).toHaveBeenCalledWith('org-1', 'user-2', 2, 'admin-1', db)
     expect(auditMock.recordAudit).toHaveBeenCalledTimes(1)
     await expect(res.json()).resolves.toEqual({
       success: true,
@@ -199,9 +201,10 @@ describe('PUT /api/organizations/[id]/members/[memberId]/usage-limit', () => {
   })
 
   it('clears the cap when creditLimit is null', async () => {
+    queueTableRows(member, [{ role: 'admin' }])
     const res = await PUT(putRequest({ creditLimit: null }), context())
     expect(res.status).toBe(200)
-    expect(mockSetOrgMemberUsageLimit).toHaveBeenCalledWith('org-1', 'user-2', null, 'admin-1')
+    expect(mockSetOrgMemberUsageLimit).toHaveBeenCalledWith('org-1', 'user-2', null, 'admin-1', db)
   })
 
   it.each([400, null])(

apps/sim/app/api/v2/organizations/[organizationId]/usage/route.test.ts

@@ -2,6 +2,7 @@
 
 import { recordAudit } from '@sim/audit'
 import type { OAuthAccessTokenPrincipal, Principal } from '@sim/auth/principal'
+import { db } from '@sim/db'
 import { member } from '@sim/db/schema'
 import { queueTableRows, resetDbChainMock, resetEnvFlagsMock, setEnvFlags } from '@sim/testing'
 import { NextRequest } from 'next/server'
@@ -144,13 +145,14 @@ describe('organization credit-limit API', () => {
     async (principal) => {
       authenticate(principal)
       admin()
+      admin()
       const response = await setLimit(
         request('members/external-user/usage-limit', { creditLimit: 400 }),
         context
       )
       expect(response.status).toBe(200)
       expect(await response.json()).toEqual({ data: { creditLimit: 400 } })
-      expect(mocks.setLimit).toHaveBeenCalledWith('org', 'external-user', 2, 'actor')
+      expect(mocks.setLimit).toHaveBeenCalledWith('org', 'external-user', 2, 'actor', db)
       expect(mocks.limitTarget).toHaveBeenCalledWith('org', 'external-user')
       expect(recordAudit).toHaveBeenCalledExactlyOnceWith(
         expect.objectContaining({
@@ -167,14 +169,69 @@ describe('organization credit-limit API', () => {
   )
 
   it.each([null, 0])('supports the distinct cap value %s', async (creditLimit) => {
+    admin()
     admin()
     expect(
       (await setLimit(request('members/external-user/usage-limit', { creditLimit }), context))
         .status
     ).toBe(200)
-    expect(mocks.setLimit).toHaveBeenCalledWith('org', 'external-user', creditLimit, 'actor')
+    expect(mocks.setLimit).toHaveBeenCalledWith('org', 'external-user', creditLimit, 'actor', db)
+  })
+
+  it('rechecks the target after acquiring mutation locks', async () => {
+    admin()
+    admin()
+    mocks.limitTarget.mockResolvedValueOnce(true).mockResolvedValueOnce(false)
+    const response = await setLimit(
+      request('members/external-user/usage-limit', { creditLimit: 400 }),
+      context
+    )
+    expect(response.status).toBe(404)
+    expect(mocks.limitTarget).toHaveBeenLastCalledWith('org', 'external-user', {
+      executor: db,
+      forShare: true,
+    })
+    expect(mocks.setLimit).not.toHaveBeenCalled()
+    expect(recordAudit).not.toHaveBeenCalled()
+  })
+
+  it('refuses an actor demoted while waiting for mutation locks', async () => {
+    admin()
+    queueTableRows(member, [{ role: 'member' }])
+    const response = await setLimit(
+      request('members/external-user/usage-limit', { creditLimit: 400 }),
+      context
+    )
+    expect(response.status).toBe(403)
+    expect(mocks.setLimit).not.toHaveBeenCalled()
+    expect(recordAudit).not.toHaveBeenCalled()
   })
 
+  it.each([
+    [personal, { disablePersonalApiKeys: true }],
+    [oauth, { disableOAuthAppAccess: true }],
+    [{ ...oauth, clientId: SIM_CLI_CLIENT_ID }, { disableCliAccess: true }],
+  ] as const)(
+    'rechecks $0.kind credential policy inside the mutation',
+    async (principal, restriction) => {
+      authenticate(principal)
+      admin()
+      admin()
+      mocks.config.mockResolvedValueOnce(null).mockResolvedValueOnce({
+        ...DEFAULT_PERMISSION_GROUP_CONFIG,
+        ...restriction,
+      })
+      const response = await setLimit(
+        request('members/external-user/usage-limit', { creditLimit: 400 }),
+        context
+      )
+      expect(response.status).toBe(403)
+      expect(mocks.config).toHaveBeenLastCalledWith('org', db)
+      expect(mocks.setLimit).not.toHaveBeenCalled()
+      expect(recordAudit).not.toHaveBeenCalled()
+    }
+  )
+
   it('returns credits and the resolved organization billing interval', async () => {
     admin()
     const response = await getLimit(request('members/external-user/usage-limit'), context)

apps/sim/lib/billing/application/member-usage-limits/use-cases.ts

@@ -1,4 +1,5 @@
 import { AuditAction, AuditResourceType } from '@sim/audit'
+import { db } from '@sim/db'
 import { createLogger } from '@sim/logger'
 import { memberUsageLimitOperations } from '@/lib/billing/application/member-usage-limits/operations'
 import { getOrganizationSubscription } from '@/lib/billing/core/billing'
@@ -10,10 +11,12 @@ import {
   isOrgMemberUsageLimitTarget,
   setOrgMemberUsageLimit,
 } from '@/lib/billing/organizations/member-limits'
+import { acquireOrganizationUserMutationLocks } from '@/lib/billing/organizations/membership'
 import {
   defineAuthorizedOrganizationUseCase,
   type OrganizationUseCaseContext,
 } from '@/lib/core/application/authorized-organization-use-case'
+import { authorizeOrganizationOperation } from '@/lib/core/application/organization-authorization'
 import { isHosted } from '@/lib/core/config/env-flags'
 import { OrchestrationError } from '@/lib/core/orchestration/types'
 
@@ -67,16 +70,33 @@ export const updateOrganizationMemberUsageLimit = defineAuthorizedOrganizationUs
   operation: memberUsageLimitOperations.update,
   authorizeResource: requireMemberUsageLimitTarget,
   async execute({
+    principal,
     input,
     context,
   }: OrganizationUseCaseContext<UpdateOrganizationMemberUsageLimitInput>) {
     const { organizationId, userId, creditLimit } = input
-    await setOrgMemberUsageLimit(
-      organizationId,
-      userId,
-      creditLimit === null ? null : creditsToDollars(creditLimit),
-      context.userId
-    )
+    await db.transaction(async (tx) => {
+      await acquireOrganizationUserMutationLocks(tx, { userId, organizationIds: [organizationId] })
+      await authorizeOrganizationOperation(principal, memberUsageLimitOperations.update, input, {
+        executor: tx,
+        forUpdate: true,
+      })
+      if (
+        !(await isOrgMemberUsageLimitTarget(organizationId, userId, {
+          executor: tx,
+          forShare: true,
+        }))
+      ) {
+        throw new OrchestrationError('not_found', 'Member not found')
+      }
+      await setOrgMemberUsageLimit(
+        organizationId,
+        userId,
+        creditLimit === null ? null : creditsToDollars(creditLimit),
+        context.userId,
+        tx
+      )
+    })
     logger.info('Updated per-member usage limit', {
       organizationId,
       memberId: userId,

apps/sim/lib/billing/organizations/member-limits.postgres.test.ts

@@ -1,45 +1,63 @@
 /** @vitest-environment node */
 
+import { recordAudit } from '@sim/audit'
+import * as schema from '@sim/db/schema'
+import { resetEnvFlagsMock, setEnvFlags } from '@sim/testing'
 import { generateId } from '@sim/utils/id'
+import { eq } from 'drizzle-orm'
 import { drizzle } from 'drizzle-orm/postgres-js'
 import postgres from 'postgres'
 import { afterAll, beforeAll, describe, expect, it, vi } from 'vitest'
 
-const { databaseUrl, select } = vi.hoisted(() => {
+const { databaseUrl, select, transaction } = vi.hoisted(() => {
   const databaseUrl = process.env.BILLING_USAGE_TEST_DATABASE_URL
   if (databaseUrl && !['localhost', '127.0.0.1', '[::1]'].includes(new URL(databaseUrl).hostname)) {
     throw new Error('Usage integration tests require a disposable local database')
   }
-  return { databaseUrl, select: vi.fn() }
+  return { databaseUrl, select: vi.fn(), transaction: vi.fn() }
 })
 vi.unmock('drizzle-orm')
 vi.unmock('@sim/db/schema')
-vi.mock('@sim/db', () => ({ db: { select }, dbReplica: { select } }))
+vi.mock('@sim/db', () => ({ db: { select, transaction }, dbReplica: { select } }))
+vi.mock('@sim/audit', async (original) => ({
+  ...(await original<typeof import('@sim/audit')>()),
+  recordAudit: vi.fn(),
+}))
 vi.mock('@/lib/billing/core/billing', () => ({ getOrganizationSubscription: vi.fn() }))
 
+import { updateOrganizationMemberUsageLimit } from '@/lib/billing/application/member-usage-limits/use-cases'
 import { isOrgMemberUsageLimitTarget } from '@/lib/billing/organizations/member-limits'
+import { acquireOrganizationUserMutationLocks } from '@/lib/billing/organizations/membership'
 
 const schemaName = `member_limits_${generateId().replaceAll('-', '')}`
 const connection = databaseUrl
   ? postgres(databaseUrl, {
-      max: 1,
+      max: 3,
       prepare: false,
-      connection: { search_path: schemaName },
+      connection: { search_path: schemaName, application_name: schemaName },
       onnotice: () => undefined,
     })
   : undefined
 
+const database = connection ? drizzle(connection, { schema }) : undefined
+
 beforeAll(async () => {
   if (!connection) return
   await connection.unsafe(`CREATE SCHEMA "${schemaName}"`)
   await connection.unsafe(`
-    CREATE TABLE member (id text PRIMARY KEY, organization_id text, user_id text);
+    CREATE TABLE member (id text PRIMARY KEY, organization_id text, user_id text, role text DEFAULT 'member');
     CREATE TABLE workspace (id text PRIMARY KEY, organization_id text, archived_at timestamp);
     CREATE TABLE permissions (id text PRIMARY KEY, user_id text, entity_type text, entity_id text);
     CREATE TABLE "user" (id text PRIMARY KEY);
     INSERT INTO "user" VALUES ('unrelated');
-    INSERT INTO member VALUES
+    INSERT INTO member (id, organization_id, user_id) VALUES
       ('m1', 'org', 'member'), ('m2', 'other', 'external'), ('m3', 'other', 'foreign-member');
+    INSERT INTO member VALUES ('actor', 'org', 'actor', 'admin');
+    CREATE TABLE organization_member_usage_limit (
+      id text PRIMARY KEY, organization_id text, user_id text, usage_limit numeric,
+      set_by text, created_at timestamp DEFAULT now(), updated_at timestamp DEFAULT now(),
+      UNIQUE(organization_id, user_id)
+    );
     INSERT INTO workspace VALUES
       ('local', 'org', null), ('foreign', 'other', null), ('archived', 'org', now());
     INSERT INTO permissions VALUES
@@ -50,11 +68,13 @@ beforeAll(async () => {
       ('p5', 'missing-workspace', 'workspace', 'missing'),
       ('p6', 'revoked', 'workspace', 'local');
   `)
-  const database = drizzle(connection)
-  select.mockImplementation((fields) => database.select(fields))
+  select.mockImplementation((fields) => database!.select(fields))
+  transaction.mockImplementation((callback) => database!.transaction(callback))
+  setEnvFlags({ isHosted: true })
 })
 
 afterAll(async () => {
+  resetEnvFlagsMock()
   if (!connection) return
   await connection.unsafe(`DROP SCHEMA "${schemaName}" CASCADE`)
   await connection.end()
@@ -81,3 +101,78 @@ describe.skipIf(!databaseUrl)('organization credit-limit target SQL', () => {
     expect(await isOrgMemberUsageLimitTarget('org', 'revoked')).toBe(false)
   })
 })
+
+describe.skipIf(!databaseUrl)('organization credit-limit mutation races', () => {
+  it.each(['member', 'external', 'archived-external'])(
+    'sets and clears a cap for the eligible target %s',
+    async (userId) => {
+      const principal = { kind: 'session', userId: 'actor', sessionId: 'session' } as const
+      await expect(
+        updateOrganizationMemberUsageLimit.execute({
+          principal,
+          input: { organizationId: 'org', userId, creditLimit: 400 },
+        })
+      ).resolves.toEqual({ creditLimit: 400 })
+      const [cap] =
+        await connection!`SELECT usage_limit, set_by FROM organization_member_usage_limit WHERE user_id = ${userId}`
+      expect(Number(cap.usage_limit)).toBe(2)
+      expect(cap.set_by).toBe('actor')
+      await expect(
+        updateOrganizationMemberUsageLimit.execute({
+          principal,
+          input: { organizationId: 'org', userId, creditLimit: null },
+        })
+      ).resolves.toEqual({ creditLimit: null })
+      expect(
+        await connection!`SELECT id FROM organization_member_usage_limit WHERE user_id = ${userId}`
+      ).toHaveLength(0)
+    }
+  )
+
+  it.each(['organization-removal', 'workspace-revocation'] as const)(
+    'rejects a target revoked by %s while the update waits',
+    async (removalKind) => {
+      const userId = `target-${removalKind}`
+      await connection!`INSERT INTO permissions VALUES (${userId}, ${userId}, 'workspace', 'local')`
+      vi.mocked(recordAudit).mockClear()
+      const ready = Promise.withResolvers<void>()
+      const release = Promise.withResolvers<void>()
+      const removal = database!.transaction(async (tx) => {
+        if (removalKind === 'organization-removal') {
+          await acquireOrganizationUserMutationLocks(tx, { userId, organizationIds: ['org'] })
+        }
+        await tx.delete(schema.permissions).where(eq(schema.permissions.userId, userId))
+        ready.resolve()
+        await release.promise
+      })
+      await ready.promise
+      const update = updateOrganizationMemberUsageLimit
+        .execute({
+          principal: { kind: 'session', userId: 'actor', sessionId: 'session' },
+          input: { organizationId: 'org', userId, creditLimit: 400 },
+        })
+        .then(
+          (result) => ({ result }),
+          (error: unknown) => ({ error })
+        )
+      try {
+        await vi.waitFor(
+          async () => {
+            const [waiting] = await connection!`SELECT count(*)::int AS count FROM pg_stat_activity
+            WHERE application_name = ${schemaName} AND wait_event_type = 'Lock'`
+            expect(waiting.count).toBeGreaterThan(0)
+          },
+          { timeout: 2000 }
+        )
+      } finally {
+        release.resolve()
+        await removal
+      }
+      expect(await update).toMatchObject({ error: { code: 'not_found' } })
+      const caps =
+        await connection!`SELECT id FROM organization_member_usage_limit WHERE user_id = ${userId}`
+      expect(caps).toHaveLength(0)
+      expect(recordAudit).not.toHaveBeenCalled()
+    }
+  )
+})

apps/sim/lib/billing/organizations/member-limits.ts

@@ -21,19 +21,24 @@ const logger = createLogger('OrgMemberLimits')
 /**
  * Includes external collaborators whose explicit workspace access belongs to the organization.
  * Retained grants on archived workspaces still qualify so their caps remain manageable.
+ * Mutations hold the organization fence before requesting a shared relationship lock;
+ * together these stabilize workspace scope and access through the write.
  */
 export async function isOrgMemberUsageLimitTarget(
   organizationId: string,
-  userId: string
+  userId: string,
+  options: { executor?: DbOrTx; forShare?: boolean } = {}
 ): Promise<boolean> {
-  const [organizationMember] = await db
+  const executor = options.executor ?? db
+  const memberQuery = executor
     .select({ id: member.id })
     .from(member)
     .where(and(eq(member.organizationId, organizationId), eq(member.userId, userId)))
     .limit(1)
+  const [organizationMember] = options.forShare ? await memberQuery.for('share') : await memberQuery
   if (organizationMember) return true
 
-  const [workspaceMember] = await db
+  const workspaceQuery = executor
     .select({ id: permissions.id })
     .from(permissions)
     .innerJoin(workspace, eq(workspace.id, permissions.entityId))
@@ -45,6 +50,9 @@ export async function isOrgMemberUsageLimitTarget(
       )
     )
     .limit(1)
+  const [workspaceMember] = options.forShare
+    ? await workspaceQuery.for('share', { of: permissions })
+    : await workspaceQuery
   return Boolean(workspaceMember)
 }