chore: import upstream snapshot with attribution

apps/sim/lib/billing/calculations/usage-monitor.test.ts

@@ -0,0 +1,129 @@
+/**
+ * @vitest-environment node
+ */
+import { beforeEach, describe, expect, it, vi } from 'vitest'
+
+const { mockFlags, mockDbLimit, mockGetOrgMemberUsageLimit, mockGetOrgMemberWorkspaceUsage } =
+  vi.hoisted(() => ({
+    mockFlags: { isHosted: true, isBillingEnabled: true },
+    mockDbLimit: vi.fn(),
+    mockGetOrgMemberUsageLimit: vi.fn(),
+    mockGetOrgMemberWorkspaceUsage: vi.fn(),
+  }))
+
+vi.mock('@/lib/core/config/env-flags', () => ({
+  get isHosted() {
+    return mockFlags.isHosted
+  },
+  get isBillingEnabled() {
+    return mockFlags.isBillingEnabled
+  },
+}))
+
+vi.mock('@sim/db', () => ({
+  db: {
+    select: () => ({
+      from: () => ({
+        where: () => ({
+          limit: mockDbLimit,
+        }),
+      }),
+    }),
+  },
+}))
+
+vi.mock('@/lib/billing/organizations/member-limits', () => ({
+  getOrgMemberUsageLimit: mockGetOrgMemberUsageLimit,
+  getOrgMemberWorkspaceUsage: mockGetOrgMemberWorkspaceUsage,
+}))
+
+// core/usage pulls in the email-rendering chain at import; stub the two symbols
+// usage-monitor imports from it so the module loads in a node test env.
+vi.mock('@/lib/billing/core/usage', () => ({
+  getPooledOrgCurrentPeriodCost: vi.fn(),
+  getUserUsageLimit: vi.fn(),
+}))
+
+import { checkOrgMemberUsageLimit } from '@/lib/billing/calculations/usage-monitor'
+
+describe('checkOrgMemberUsageLimit', () => {
+  beforeEach(() => {
+    vi.clearAllMocks()
+    mockFlags.isHosted = true
+    mockFlags.isBillingEnabled = true
+    mockDbLimit.mockResolvedValue([{ organizationId: 'org-1' }])
+    mockGetOrgMemberUsageLimit.mockResolvedValue(2)
+    mockGetOrgMemberWorkspaceUsage.mockResolvedValue(1)
+  })
+
+  it('no-ops when not hosted', async () => {
+    mockFlags.isHosted = false
+    const result = await checkOrgMemberUsageLimit('user-1', 'ws-1')
+    expect(result.isExceeded).toBe(false)
+    expect(mockDbLimit).not.toHaveBeenCalled()
+    expect(mockGetOrgMemberUsageLimit).not.toHaveBeenCalled()
+  })
+
+  it('no-ops when billing is disabled', async () => {
+    mockFlags.isBillingEnabled = false
+    const result = await checkOrgMemberUsageLimit('user-1', 'ws-1')
+    expect(result.isExceeded).toBe(false)
+  })
+
+  it('no-ops when workspaceId is empty', async () => {
+    const result = await checkOrgMemberUsageLimit('user-1', '')
+    expect(result.isExceeded).toBe(false)
+    expect(mockDbLimit).not.toHaveBeenCalled()
+  })
+
+  it('no-ops when the workspace is not org-owned', async () => {
+    mockDbLimit.mockResolvedValue([{ organizationId: null }])
+    const result = await checkOrgMemberUsageLimit('user-1', 'ws-1')
+    expect(result.isExceeded).toBe(false)
+    expect(mockGetOrgMemberUsageLimit).not.toHaveBeenCalled()
+  })
+
+  it('no-ops when the member has no cap set', async () => {
+    mockGetOrgMemberUsageLimit.mockResolvedValue(null)
+    const result = await checkOrgMemberUsageLimit('user-1', 'ws-1')
+    expect(result.isExceeded).toBe(false)
+    expect(mockGetOrgMemberWorkspaceUsage).not.toHaveBeenCalled()
+  })
+
+  it('does not block when usage is below the cap', async () => {
+    mockGetOrgMemberWorkspaceUsage.mockResolvedValue(1)
+    mockGetOrgMemberUsageLimit.mockResolvedValue(2)
+    const result = await checkOrgMemberUsageLimit('user-1', 'ws-1')
+    expect(result.isExceeded).toBe(false)
+    expect(result.currentUsage).toBe(1)
+    expect(result.limit).toBe(2)
+    expect(result.message).toBeUndefined()
+  })
+
+  it('blocks when usage meets the cap (>=)', async () => {
+    mockGetOrgMemberWorkspaceUsage.mockResolvedValue(2)
+    mockGetOrgMemberUsageLimit.mockResolvedValue(2)
+    const result = await checkOrgMemberUsageLimit('user-1', 'ws-1')
+    expect(result.isExceeded).toBe(true)
+    expect(result.message).toBeTruthy()
+  })
+
+  it('blocks all usage when the cap is 0', async () => {
+    mockGetOrgMemberUsageLimit.mockResolvedValue(0)
+    mockGetOrgMemberWorkspaceUsage.mockResolvedValue(0)
+    const result = await checkOrgMemberUsageLimit('user-1', 'ws-1')
+    expect(result.isExceeded).toBe(true)
+  })
+
+  it('fails open when an unexpected error occurs', async () => {
+    mockDbLimit.mockRejectedValue(new Error('db down'))
+    const result = await checkOrgMemberUsageLimit('user-1', 'ws-1')
+    expect(result.isExceeded).toBe(false)
+  })
+
+  it('fails open when org-workspace usage cannot be computed (unexpected error)', async () => {
+    mockGetOrgMemberWorkspaceUsage.mockRejectedValue(new Error('db unavailable'))
+    const result = await checkOrgMemberUsageLimit('user-1', 'ws-1')
+    expect(result.isExceeded).toBe(false)
+  })
+})

apps/sim/lib/billing/calculations/usage-monitor.ts

@@ -0,0 +1,526 @@
+import { db } from '@sim/db'
+import { member, userStats, workspace } from '@sim/db/schema'
+import { createLogger } from '@sim/logger'
+import { toError } from '@sim/utils/errors'
+import { and, eq } from 'drizzle-orm'
+import { defaultBillingPeriod } from '@/lib/billing/core/billing-period'
+import {
+  getHighestPrioritySubscription,
+  type HighestPrioritySubscription,
+} from '@/lib/billing/core/plan'
+import { getPooledOrgCurrentPeriodCost, getUserUsageLimit } from '@/lib/billing/core/usage'
+import { getBillingPeriodUsageCost } from '@/lib/billing/core/usage-log'
+import { dollarsToCredits } from '@/lib/billing/credits/conversion'
+import {
+  computeDailyRefreshConsumed,
+  getOrgMemberRefreshBounds,
+} from '@/lib/billing/credits/daily-refresh'
+import {
+  getOrgMemberUsageLimit,
+  getOrgMemberWorkspaceUsage,
+} from '@/lib/billing/organizations/member-limits'
+import { getPlanTierDollars, isPaid } from '@/lib/billing/plan-helpers'
+import { isOrgScopedSubscription } from '@/lib/billing/subscriptions/utils'
+import { toDecimal, toNumber } from '@/lib/billing/utils/decimal'
+import { isBillingEnabled, isHosted } from '@/lib/core/config/env-flags'
+
+const logger = createLogger('UsageMonitor')
+
+const WARNING_THRESHOLD = 80
+
+interface UsageData {
+  percentUsed: number
+  isWarning: boolean
+  isExceeded: boolean
+  currentUsage: number
+  limit: number
+  /**
+   * Whether the returned values are this user's individual slice or the
+   * organization's pooled total/cap. When an org pool is the blocker,
+   * the pooled values are surfaced here so error messages reflect it.
+   */
+  scope: 'user' | 'organization'
+  /** Present only when `scope === 'organization'`. */
+  organizationId: string | null
+}
+
+async function computePooledOrgUsage(
+  organizationId: string,
+  sub: {
+    plan: string | null
+    seats: number | null
+    periodStart: Date | null
+    periodEnd: Date | null
+  }
+): Promise<number> {
+  const { memberIds, currentPeriodCost } = await getPooledOrgCurrentPeriodCost(organizationId)
+  if (memberIds.length === 0) return 0
+
+  const billingPeriod =
+    sub.periodStart && sub.periodEnd
+      ? { start: sub.periodStart, end: sub.periodEnd }
+      : defaultBillingPeriod()
+  const ledgerUsage = await getBillingPeriodUsageCost(
+    { type: 'organization', id: organizationId },
+    billingPeriod
+  )
+
+  return applyOrgRefresh(organizationId, sub, currentPeriodCost + ledgerUsage, memberIds)
+}
+
+/**
+ * Checks a user's cost usage against their subscription plan limit
+ * and returns usage information including whether they're approaching the limit
+ */
+export async function checkUsageStatus(
+  userId: string,
+  preloadedSubscription?: HighestPrioritySubscription
+): Promise<UsageData> {
+  try {
+    if (!isBillingEnabled) {
+      const statsRecords = await db.select().from(userStats).where(eq(userStats.userId, userId))
+      const currentUsage =
+        statsRecords.length > 0 ? toNumber(toDecimal(statsRecords[0].currentPeriodCost)) : 0
+
+      return {
+        percentUsed: Math.min((currentUsage / 1000) * 100, 100),
+        isWarning: false,
+        isExceeded: false,
+        currentUsage,
+        limit: 1000,
+        scope: 'user',
+        organizationId: null,
+      }
+    }
+
+    const sub =
+      preloadedSubscription !== undefined
+        ? preloadedSubscription
+        : await getHighestPrioritySubscription(userId)
+
+    const limit = await getUserUsageLimit(userId, sub)
+    logger.info('Using stored usage limit', { userId, limit })
+
+    const subIsOrgScoped = isOrgScopedSubscription(sub, userId)
+    const scope: 'user' | 'organization' = subIsOrgScoped ? 'organization' : 'user'
+    const organizationId: string | null = subIsOrgScoped && sub ? sub.referenceId : null
+
+    if (subIsOrgScoped && sub) {
+      const currentUsage = await computePooledOrgUsage(sub.referenceId, sub)
+      return buildUsageData({ currentUsage, limit, scope, organizationId })
+    }
+
+    const statsRecords = await db
+      .select()
+      .from(userStats)
+      .where(eq(userStats.userId, userId))
+      .limit(1)
+
+    if (statsRecords.length === 0) {
+      logger.info('No usage stats found for user', { userId, limit })
+      return {
+        percentUsed: 0,
+        isWarning: false,
+        isExceeded: false,
+        currentUsage: 0,
+        limit,
+        scope: 'user',
+        organizationId: null,
+      }
+    }
+
+    const billingPeriod =
+      sub?.periodStart && sub.periodEnd
+        ? { start: sub.periodStart, end: sub.periodEnd }
+        : defaultBillingPeriod()
+    const ledgerUsage = await getBillingPeriodUsageCost({ type: 'user', id: userId }, billingPeriod)
+    let currentUsage = toNumber(toDecimal(statsRecords[0].currentPeriodCost)) + ledgerUsage
+    if (sub && isPaid(sub.plan) && sub.periodStart) {
+      const planDollars = getPlanTierDollars(sub.plan)
+      if (planDollars > 0) {
+        const refresh = await computeDailyRefreshConsumed({
+          userIds: [userId],
+          periodStart: sub.periodStart,
+          periodEnd: sub.periodEnd ?? null,
+          planDollars,
+          billingEntity: { type: 'user', id: userId },
+        })
+        currentUsage = Math.max(0, currentUsage - refresh)
+      }
+    }
+
+    return buildUsageData({ currentUsage, limit, scope, organizationId })
+  } catch (error) {
+    logger.error('Error checking usage status', {
+      error: error instanceof Error ? { message: error.message, stack: error.stack } : error,
+      userId,
+    })
+
+    logger.error('Cannot determine usage status - blocking execution', {
+      userId,
+      error: toError(error).message,
+    })
+
+    return {
+      percentUsed: 100,
+      isWarning: false,
+      isExceeded: true,
+      currentUsage: 0,
+      limit: 0,
+      scope: 'user',
+      organizationId: null,
+    }
+  }
+}
+
+async function applyOrgRefresh(
+  organizationId: string,
+  sub: {
+    plan: string | null
+    seats: number | null
+    periodStart: Date | null
+    periodEnd: Date | null
+  },
+  currentUsage: number,
+  preloadedMemberIds?: string[]
+): Promise<number> {
+  if (!isPaid(sub.plan) || !sub.periodStart) {
+    return currentUsage
+  }
+
+  const memberIds =
+    preloadedMemberIds ?? (await getPooledOrgCurrentPeriodCost(organizationId)).memberIds
+  if (memberIds.length === 0) return currentUsage
+
+  const planDollars = getPlanTierDollars(sub.plan)
+  if (planDollars <= 0) return currentUsage
+
+  const userBounds = await getOrgMemberRefreshBounds(organizationId, sub.periodStart)
+  const refresh = await computeDailyRefreshConsumed({
+    userIds: memberIds,
+    periodStart: sub.periodStart,
+    periodEnd: sub.periodEnd ?? null,
+    planDollars,
+    seats: sub.seats || 1,
+    userBounds: Object.keys(userBounds).length > 0 ? userBounds : undefined,
+    billingEntity: { type: 'organization', id: organizationId },
+  })
+
+  return Math.max(0, currentUsage - refresh)
+}
+
+function buildUsageData(params: {
+  currentUsage: number
+  limit: number
+  scope: 'user' | 'organization'
+  organizationId: string | null
+}): UsageData {
+  const { currentUsage, limit, scope, organizationId } = params
+  const percentUsed = limit > 0 ? Math.min((currentUsage / limit) * 100, 100) : 100
+  const isExceeded = currentUsage >= limit
+  const isWarning = !isExceeded && percentUsed >= WARNING_THRESHOLD
+
+  logger.info('Final usage statistics', {
+    currentUsage,
+    limit,
+    percentUsed,
+    isWarning,
+    isExceeded,
+    scope,
+    organizationId,
+  })
+
+  return {
+    percentUsed,
+    isWarning,
+    isExceeded,
+    currentUsage,
+    limit,
+    scope,
+    organizationId,
+  }
+}
+
+/**
+ * Displays a notification to the user when they're approaching their usage limit
+ * Can be called on app startup or before executing actions that might incur costs
+ */
+async function checkAndNotifyUsage(userId: string): Promise<void> {
+  try {
+    if (!isBillingEnabled) {
+      return
+    }
+
+    const usageData = await checkUsageStatus(userId)
+
+    if (usageData.isExceeded) {
+      logger.warn('User has exceeded usage limits', {
+        userId,
+        usage: usageData.currentUsage,
+        limit: usageData.limit,
+      })
+
+      if (typeof window !== 'undefined') {
+        window.dispatchEvent(
+          new CustomEvent('usage-exceeded', {
+            detail: { usageData },
+          })
+        )
+      }
+    } else if (usageData.isWarning) {
+      logger.info('User approaching usage limits', {
+        userId,
+        usage: usageData.currentUsage,
+        limit: usageData.limit,
+        percent: usageData.percentUsed,
+      })
+
+      if (typeof window !== 'undefined') {
+        window.dispatchEvent(
+          new CustomEvent('usage-warning', {
+            detail: { usageData },
+          })
+        )
+      }
+    }
+  } catch (error) {
+    logger.error('Error in usage notification system', { error, userId })
+  }
+}
+
+/**
+ * Whether an account (or its organization owner) is billing-blocked — frozen for
+ * a dispute or flagged for a payment issue. Independent of usage limits, so it can
+ * gate paths that are otherwise exempt from metered caps (e.g. BYOK wand): a
+ * frozen account is locked out everywhere.
+ */
+export async function checkBillingBlocked(
+  userId: string
+): Promise<{ blocked: boolean; message?: string }> {
+  const stats = await db
+    .select({ blocked: userStats.billingBlocked, blockedReason: userStats.billingBlockedReason })
+    .from(userStats)
+    .where(eq(userStats.userId, userId))
+    .limit(1)
+
+  if (stats.length > 0 && stats[0].blocked) {
+    return {
+      blocked: true,
+      message:
+        stats[0].blockedReason === 'dispute'
+          ? 'Account frozen. Please contact support to resolve this issue.'
+          : 'Billing issue detected. Please update your payment method to continue.',
+    }
+  }
+
+  const memberships = await db
+    .select({ organizationId: member.organizationId })
+    .from(member)
+    .where(eq(member.userId, userId))
+
+  for (const m of memberships) {
+    const owners = await db
+      .select({ userId: member.userId })
+      .from(member)
+      .where(and(eq(member.organizationId, m.organizationId), eq(member.role, 'owner')))
+      .limit(1)
+
+    if (owners.length > 0) {
+      const ownerStats = await db
+        .select({
+          blocked: userStats.billingBlocked,
+          blockedReason: userStats.billingBlockedReason,
+        })
+        .from(userStats)
+        .where(eq(userStats.userId, owners[0].userId))
+        .limit(1)
+
+      if (ownerStats.length > 0 && ownerStats[0].blocked) {
+        return {
+          blocked: true,
+          message:
+            ownerStats[0].blockedReason === 'dispute'
+              ? 'Organization account frozen. Please contact support to resolve this issue.'
+              : 'Organization billing issue. Please contact your organization owner.',
+        }
+      }
+    }
+  }
+
+  return { blocked: false }
+}
+
+/**
+ * Server-side function to check if a user has exceeded their usage limits
+ * For use in API routes, webhooks, and scheduled executions
+ *
+ * @param userId The ID of the user to check
+ * @returns An object containing the exceeded status and usage details
+ */
+export async function checkServerSideUsageLimits(
+  userId: string,
+  preloadedSubscription?: HighestPrioritySubscription
+): Promise<{
+  isExceeded: boolean
+  currentUsage: number
+  limit: number
+  message?: string
+}> {
+  try {
+    if (!isBillingEnabled) {
+      return {
+        isExceeded: false,
+        currentUsage: 0,
+        limit: 99999,
+      }
+    }
+
+    logger.info('Server-side checking usage limits for user', { userId })
+
+    const stats = await db
+      .select({ current: userStats.currentPeriodCost })
+      .from(userStats)
+      .where(eq(userStats.userId, userId))
+      .limit(1)
+
+    const currentUsage = stats.length > 0 ? toNumber(toDecimal(stats[0].current)) : 0
+
+    const blocked = await checkBillingBlocked(userId)
+    if (blocked.blocked) {
+      return { isExceeded: true, currentUsage, limit: 0, message: blocked.message }
+    }
+
+    const usageData = await checkUsageStatus(userId, preloadedSubscription)
+
+    const formattedUsage = (usageData.currentUsage ?? 0).toFixed(2)
+    const formattedLimit = (usageData.limit ?? 0).toFixed(2)
+    const exceededMessage =
+      usageData.scope === 'organization'
+        ? `Organization usage limit exceeded: $${formattedUsage} pooled of $${formattedLimit} organization limit. Ask a team admin to raise the organization usage limit to continue.`
+        : `Usage limit exceeded: $${formattedUsage} used of $${formattedLimit} limit. Please upgrade your plan or raise your usage limit to continue.`
+
+    return {
+      isExceeded: usageData.isExceeded,
+      currentUsage: usageData.currentUsage,
+      limit: usageData.limit,
+      message: usageData.isExceeded ? exceededMessage : undefined,
+    }
+  } catch (error) {
+    logger.error('Error in server-side usage limit check', {
+      error: error instanceof Error ? { message: error.message, stack: error.stack } : error,
+      userId,
+    })
+
+    logger.error('Cannot determine usage limits - blocking execution', {
+      userId,
+      error: toError(error).message,
+    })
+
+    return {
+      isExceeded: true,
+      currentUsage: 0,
+      limit: 0,
+      message:
+        error instanceof Error && error.message.includes('No user stats record found')
+          ? 'User account not properly initialized. Please contact support.'
+          : 'Unable to determine usage limits. Execution blocked for security. Please contact support.',
+    }
+  }
+}
+
+/**
+ * Per-member usage cap for the organization that owns the given workspace.
+ *
+ * Hosted-only and independent of the pooled org limit
+ * ({@link checkServerSideUsageLimits}). No-ops (isExceeded:false) when the
+ * feature is off (`!isHosted` / `!isBillingEnabled`), the workspace is not
+ * org-owned, or the actor has no per-member cap. Otherwise compares the actor's
+ * current-period usage inside the org's workspaces against their cap
+ * (`usage >= limit` blocks).
+ *
+ * Fails open on unexpected error: this is a secondary, additive gate, so a
+ * transient fault must not block execution that the primary pooled/personal
+ * check already allowed.
+ */
+export async function checkOrgMemberUsageLimit(
+  userId: string,
+  workspaceId: string
+): Promise<{
+  isExceeded: boolean
+  currentUsage: number
+  limit: number | null
+  message?: string
+}> {
+  try {
+    if (!isHosted || !isBillingEnabled || !workspaceId) {
+      return { isExceeded: false, currentUsage: 0, limit: null }
+    }
+
+    const [workspaceRow] = await db
+      .select({ organizationId: workspace.organizationId })
+      .from(workspace)
+      .where(eq(workspace.id, workspaceId))
+      .limit(1)
+
+    const organizationId = workspaceRow?.organizationId
+    if (!organizationId) {
+      return { isExceeded: false, currentUsage: 0, limit: null }
+    }
+
+    // Resolve the cap first and short-circuit when unset (the common case); only
+    // then is computing usage worthwhile. Kept sequential, not raced, to avoid a
+    // usage query on every uncapped member's execution.
+    const limit = await getOrgMemberUsageLimit(organizationId, userId)
+    if (limit === null) {
+      return { isExceeded: false, currentUsage: 0, limit: null }
+    }
+
+    const usage = await getOrgMemberWorkspaceUsage(organizationId, userId)
+    const isExceeded = usage >= limit
+
+    return {
+      isExceeded,
+      currentUsage: usage,
+      limit,
+      message: isExceeded
+        ? `Member credit limit exceeded: ${dollarsToCredits(usage).toLocaleString()} of ${dollarsToCredits(limit).toLocaleString()} credits used for this organization's workspaces. Ask an organization admin to raise your credit limit to continue.`
+        : undefined,
+    }
+  } catch (error) {
+    logger.error('Error checking per-member org usage limit', {
+      error: toError(error).message,
+      userId,
+      workspaceId,
+    })
+    return { isExceeded: false, currentUsage: 0, limit: null }
+  }
+}
+
+/**
+ * Unified usage gate for an actor: the pooled/personal cap first
+ * ({@link checkServerSideUsageLimits}), then the per-member org-workspace cap
+ * ({@link checkOrgMemberUsageLimit}) when a workspace is in scope. Returns the
+ * first exceeded result so every billable surface (workflow exec, copilot,
+ * voice, wand, enrichment, KB indexing) can gate on a single
+ * `{ isExceeded, message }`. `scope` distinguishes a pooled cap from a per-member
+ * cap so clients can hide an "Upgrade" affordance a capped member can't act on (a
+ * per-member cap is only raisable by an org admin).
+ */
+export async function checkActorUsageLimits(
+  userId: string,
+  workspaceId?: string | null
+): Promise<{ isExceeded: boolean; message?: string; scope?: 'pooled' | 'member' }> {
+  const pooled = await checkServerSideUsageLimits(userId)
+  if (pooled.isExceeded) {
+    return { isExceeded: true, message: pooled.message, scope: 'pooled' }
+  }
+
+  if (workspaceId) {
+    const member = await checkOrgMemberUsageLimit(userId, workspaceId)
+    if (member.isExceeded) {
+      return { isExceeded: true, message: member.message, scope: 'member' }
+    }
+  }
+
+  return { isExceeded: false }
+}

apps/sim/lib/billing/calculations/usage-reservation.test.ts

@@ -0,0 +1,152 @@
+/**
+ * @vitest-environment node
+ */
+import { redisConfigMock, redisConfigMockFns } from '@sim/testing'
+import { beforeEach, describe, expect, it, vi } from 'vitest'
+
+const { mockFlags } = vi.hoisted(() => ({
+  mockFlags: { isBillingEnabled: true },
+}))
+
+vi.mock('@/lib/core/config/env-flags', () => ({
+  get isBillingEnabled() {
+    return mockFlags.isBillingEnabled
+  },
+  isHosted: true,
+}))
+
+vi.mock('@/lib/core/config/redis', () => redisConfigMock)
+
+import {
+  releaseExecutionSlot,
+  reserveExecutionSlot,
+  resolveBillingEntityKey,
+} from '@/lib/billing/calculations/usage-reservation'
+
+const evalMock = vi.fn()
+const getdelMock = vi.fn()
+const zremMock = vi.fn()
+const fakeRedis = { eval: evalMock, getdel: getdelMock, zrem: zremMock }
+
+const baseParams = {
+  userId: 'user-1',
+  executionId: 'exec-1',
+  subscription: { plan: 'free' as const, referenceId: 'user-1' },
+  currentUsage: 0,
+  limit: 5,
+}
+
+describe('usage-reservation', () => {
+  beforeEach(() => {
+    vi.clearAllMocks()
+    mockFlags.isBillingEnabled = true
+    redisConfigMockFns.mockGetRedisClient.mockReturnValue(fakeRedis)
+  })
+
+  describe('resolveBillingEntityKey', () => {
+    it('keys personal subscriptions by user', () => {
+      expect(resolveBillingEntityKey('user-1', { referenceId: 'user-1' })).toBe('user:user-1')
+    })
+
+    it('keys org-scoped subscriptions by organization', () => {
+      expect(resolveBillingEntityKey('user-1', { referenceId: 'org-9' })).toBe('org:org-9')
+    })
+  })
+
+  describe('reserveExecutionSlot', () => {
+    it('admits when the reservation script returns 1', async () => {
+      evalMock.mockResolvedValueOnce(1)
+      const result = await reserveExecutionSlot(baseParams)
+      expect(result.reserved).toBe(true)
+      expect(evalMock).toHaveBeenCalledTimes(1)
+    })
+
+    it('rejects when the reservation script returns 0 (slots full)', async () => {
+      evalMock.mockResolvedValueOnce(0)
+      const result = await reserveExecutionSlot(baseParams)
+      expect(result.reserved).toBe(false)
+    })
+
+    it('passes the free-tier concurrency cap and headroom slots to the script', async () => {
+      evalMock.mockResolvedValueOnce(1)
+      await reserveExecutionSlot(baseParams)
+      const args = evalMock.mock.calls[0]
+      // eval(script, numKeys, inflightKey, pointerKey, now, expiry, maxConc, headroomSlots, member, entityKey, pttl)
+      expect(args[2]).toBe('usage:inflight:user:user-1')
+      expect(args[3]).toBe('usage:reservation:exec-1')
+      expect(args[6]).toBe('15')
+      expect(args[7]).toBe('1000')
+      expect(args[8]).toBe('exec-1')
+      expect(args[9]).toBe('user:user-1')
+    })
+
+    it('reserves against the org entity for org-scoped subscriptions', async () => {
+      evalMock.mockResolvedValueOnce(1)
+      await reserveExecutionSlot({
+        ...baseParams,
+        subscription: { plan: 'team', referenceId: 'org-9' },
+      })
+      const args = evalMock.mock.calls[0]
+      expect(args[2]).toBe('usage:inflight:org:org-9')
+      expect(args[6]).toBe('150')
+    })
+
+    it('clamps negative headroom to zero slots', async () => {
+      evalMock.mockResolvedValueOnce(0)
+      await reserveExecutionSlot({ ...baseParams, currentUsage: 10, limit: 5 })
+      expect(evalMock.mock.calls[0][7]).toBe('0')
+    })
+
+    it('fails open (admits) when billing enforcement is disabled', async () => {
+      mockFlags.isBillingEnabled = false
+      const result = await reserveExecutionSlot(baseParams)
+      expect(result.reserved).toBe(true)
+      expect(evalMock).not.toHaveBeenCalled()
+    })
+
+    it('fails open (admits) when Redis is unavailable', async () => {
+      redisConfigMockFns.mockGetRedisClient.mockReturnValue(null)
+      const result = await reserveExecutionSlot(baseParams)
+      expect(result.reserved).toBe(true)
+      expect(evalMock).not.toHaveBeenCalled()
+    })
+
+    it('fails open (admits) when the reservation script throws', async () => {
+      evalMock.mockRejectedValueOnce(new Error('connection lost'))
+      const result = await reserveExecutionSlot(baseParams)
+      expect(result.reserved).toBe(true)
+    })
+  })
+
+  describe('releaseExecutionSlot', () => {
+    it('reads the pointer then removes the slot from that entity set', async () => {
+      getdelMock.mockResolvedValueOnce('org:org-9')
+      await releaseExecutionSlot('exec-1')
+      expect(getdelMock).toHaveBeenCalledWith('usage:reservation:exec-1')
+      expect(zremMock).toHaveBeenCalledWith('usage:inflight:org:org-9', 'exec-1')
+    })
+
+    it('does not touch the in-flight set when the pointer is already gone', async () => {
+      getdelMock.mockResolvedValueOnce(null)
+      await releaseExecutionSlot('exec-1')
+      expect(zremMock).not.toHaveBeenCalled()
+    })
+
+    it('uses only single-key commands (cluster-safe; no key built inside Lua)', async () => {
+      getdelMock.mockResolvedValueOnce('user:user-1')
+      await releaseExecutionSlot('exec-1')
+      expect(evalMock).not.toHaveBeenCalled()
+    })
+
+    it('is a no-op when billing enforcement is disabled', async () => {
+      mockFlags.isBillingEnabled = false
+      await releaseExecutionSlot('exec-1')
+      expect(getdelMock).not.toHaveBeenCalled()
+    })
+
+    it('swallows release errors', async () => {
+      getdelMock.mockRejectedValueOnce(new Error('boom'))
+      await expect(releaseExecutionSlot('exec-1')).resolves.toBeUndefined()
+    })
+  })
+})

apps/sim/lib/billing/calculations/usage-reservation.ts

@@ -0,0 +1,207 @@
+import { createLogger } from '@sim/logger'
+import { toError } from '@sim/utils/errors'
+import { BASE_EXECUTION_CHARGE } from '@/lib/billing/constants'
+import { getPlanTypeForLimits } from '@/lib/billing/plan-helpers'
+import { isOrgScopedSubscription } from '@/lib/billing/subscriptions/utils'
+import { isBillingEnabled } from '@/lib/core/config/env-flags'
+import { getRedisClient } from '@/lib/core/config/redis'
+import { getExecutionReservationTtlMs } from '@/lib/core/execution-limits'
+import type { SubscriptionPlan } from '@/lib/core/rate-limiter/types'
+
+const logger = createLogger('UsageReservation')
+
+/**
+ * Maximum number of simultaneously in-flight (admitted but not-yet-costed)
+ * executions a single billing entity may hold at once.
+ *
+ * The usage-cap admission gate reads already-recorded cost, but cost is only
+ * written when an execution finishes. Without a reservation, N parallel
+ * executions all read the same pre-burst usage, all pass the cap, and all run —
+ * collectively spending far past the cap before any cost lands in the ledger
+ * (free-tier abuse / hard-cap defeat). Bounding the number of in-flight
+ * executions per billing entity bounds the worst-case overshoot to roughly this
+ * many executions' worth of spend.
+ */
+const MAX_CONCURRENT_EXECUTIONS: Record<SubscriptionPlan, number> = {
+  free: 15,
+  pro: 75,
+  team: 150,
+  enterprise: 300,
+}
+
+/**
+ * Per-slot reserved cost estimate (dollars). The guaranteed-minimum charge
+ * every execution incurs, used to taper admission as recorded usage approaches
+ * the cap: an entity may hold at most `floor(headroom / estimate)` concurrent
+ * slots, keeping `recordedUsage + reservedSlots * estimate <= limit`. A lone
+ * execution is never blocked on headroom alone — the recorded-usage gate
+ * (`isExceeded`) governs the single-execution case, so the only residual
+ * overshoot is the one already inherent to admission (cost is unknown until the
+ * execution finishes).
+ */
+const SLOT_COST_ESTIMATE = BASE_EXECUTION_CHARGE
+
+const INFLIGHT_KEY_PREFIX = 'usage:inflight:'
+const POINTER_KEY_PREFIX = 'usage:reservation:'
+
+/**
+ * Atomically admit an execution only when both the per-entity concurrency cap
+ * and the remaining usage headroom permit it, then record the in-flight slot.
+ *
+ * Prune expired members (crash safety) -> `count = ZCARD` -> reject when
+ * `count >= min(maxConcurrency, max(1, headroomSlots))` -> otherwise `ZADD` the
+ * slot, refresh the set TTL, and write the per-execution pointer for release.
+ * The `max(1, ...)` floor guarantees a lone execution is never blocked on
+ * headroom alone; concurrency above the first slot still tapers with headroom.
+ */
+const RESERVE_SCRIPT = `
+local now = tonumber(ARGV[1])
+local expiryScore = tonumber(ARGV[2])
+local maxConcurrency = tonumber(ARGV[3])
+local headroomSlots = tonumber(ARGV[4])
+local pttl = tonumber(ARGV[7])
+redis.call('ZREMRANGEBYSCORE', KEYS[1], '-inf', now)
+local count = redis.call('ZCARD', KEYS[1])
+if headroomSlots < 1 then headroomSlots = 1 end
+local allowed = maxConcurrency
+if headroomSlots < allowed then allowed = headroomSlots end
+if count >= allowed then
+  return 0
+end
+redis.call('ZADD', KEYS[1], expiryScore, ARGV[5])
+redis.call('PEXPIRE', KEYS[1], pttl)
+redis.call('SET', KEYS[2], ARGV[6], 'PX', pttl)
+return 1
+`
+
+/**
+ * Stable per-entity reservation key. Org-scoped subscriptions reserve against
+ * the organization's pooled cap; everyone else against their personal cap —
+ * mirroring the entity the usage limit itself is enforced on.
+ */
+export function resolveBillingEntityKey(
+  userId: string,
+  subscription: { referenceId?: string | null } | null | undefined
+): string {
+  if (isOrgScopedSubscription(subscription, userId) && subscription?.referenceId) {
+    return `org:${subscription.referenceId}`
+  }
+  return `user:${userId}`
+}
+
+function getMaxConcurrentExecutions(plan: string | null | undefined): number {
+  return MAX_CONCURRENT_EXECUTIONS[getPlanTypeForLimits(plan) as SubscriptionPlan]
+}
+
+export interface ReserveExecutionSlotParams {
+  userId: string
+  executionId: string
+  subscription: { plan?: string | null; referenceId?: string | null } | null | undefined
+  /** Recorded usage for the billing entity at admission time (dollars). */
+  currentUsage: number
+  /** The entity's usage cap (dollars). */
+  limit: number
+}
+
+export interface ReserveExecutionSlotResult {
+  reserved: boolean
+}
+
+/**
+ * Atomic admission reservation that closes the usage-cap check-then-use race.
+ *
+ * No-ops (admits) when billing enforcement is off or Redis is unavailable —
+ * the caller's recorded-usage check still runs in those cases, and failing open
+ * here matches the rate limiter rather than turning a Redis blip into a full
+ * execution outage.
+ */
+export async function reserveExecutionSlot(
+  params: ReserveExecutionSlotParams
+): Promise<ReserveExecutionSlotResult> {
+  if (!isBillingEnabled) {
+    return { reserved: true }
+  }
+
+  const redis = getRedisClient()
+  if (!redis) {
+    return { reserved: true }
+  }
+
+  const { userId, executionId, subscription, currentUsage, limit } = params
+  const entityKey = resolveBillingEntityKey(userId, subscription)
+  const maxConcurrency = getMaxConcurrentExecutions(subscription?.plan)
+  const headroom = Math.max(0, limit - currentUsage)
+  const headroomSlots = Math.floor(headroom / SLOT_COST_ESTIMATE)
+  const ttlMs = getExecutionReservationTtlMs()
+  const now = Date.now()
+  const expiryScore = now + ttlMs
+
+  try {
+    const result = await redis.eval(
+      RESERVE_SCRIPT,
+      2,
+      `${INFLIGHT_KEY_PREFIX}${entityKey}`,
+      `${POINTER_KEY_PREFIX}${executionId}`,
+      now.toString(),
+      expiryScore.toString(),
+      maxConcurrency.toString(),
+      headroomSlots.toString(),
+      executionId,
+      entityKey,
+      ttlMs.toString()
+    )
+
+    const reserved = result === 1
+    if (!reserved) {
+      logger.warn('Execution admission throttled — concurrency/usage reservation full', {
+        entityKey,
+        executionId,
+        maxConcurrency,
+        headroomSlots,
+      })
+    }
+    return { reserved }
+  } catch (error) {
+    logger.error('Usage reservation error — failing open (admitting execution)', {
+      error: toError(error).message,
+      entityKey,
+      executionId,
+    })
+    return { reserved: true }
+  }
+}
+
+/**
+ * Release the in-flight reservation held for an execution. Best-effort and
+ * idempotent — safe to call for executions that never reserved (Redis down,
+ * billing disabled) or are released more than once. Must NOT be called for a
+ * paused execution that may still resume.
+ *
+ * Uses discrete single-key commands rather than a Lua script that rebuilds the
+ * in-flight key from the pointer value: the entity that owns the slot is only
+ * known after reading the pointer, and constructing a key inside Lua bypasses
+ * the `KEYS` declaration that Redis Cluster relies on for slot routing.
+ */
+export async function releaseExecutionSlot(executionId: string): Promise<void> {
+  if (!isBillingEnabled) {
+    return
+  }
+
+  const redis = getRedisClient()
+  if (!redis) {
+    return
+  }
+
+  try {
+    const pointerKey = `${POINTER_KEY_PREFIX}${executionId}`
+    const entityKey = await redis.getdel(pointerKey)
+    if (entityKey) {
+      await redis.zrem(`${INFLIGHT_KEY_PREFIX}${entityKey}`, executionId)
+    }
+  } catch (error) {
+    logger.warn('Failed to release usage reservation', {
+      error: toError(error).message,
+      executionId,
+    })
+  }
+}