chore: import upstream snapshot with attribution

apps/sim/lib/billing/storage/index.ts

@@ -0,0 +1,8 @@
+export { checkStorageQuota, getUserStorageLimit, getUserStorageUsage } from './limits'
+export {
+  checkAndIncrementStorageUsageInTx,
+  decrementStorageUsage,
+  decrementStorageUsageInTx,
+  incrementStorageUsage,
+  maybeNotifyStorageLimit,
+} from './tracking'

apps/sim/lib/billing/storage/limits.ts

@@ -0,0 +1,206 @@
+/**
+ * Storage limit management
+ * Similar to cost limits but for file storage quotas
+ */
+
+import { db } from '@sim/db'
+import {
+  DEFAULT_ENTERPRISE_STORAGE_LIMIT_GB,
+  DEFAULT_FREE_STORAGE_LIMIT_GB,
+  DEFAULT_PRO_STORAGE_LIMIT_GB,
+  DEFAULT_TEAM_STORAGE_LIMIT_GB,
+} from '@sim/db/constants'
+import { organization, userStats } from '@sim/db/schema'
+import { createLogger } from '@sim/logger'
+import { eq } from 'drizzle-orm'
+import type { HighestPrioritySubscription } from '@/lib/billing/core/plan'
+import { getPlanTypeForLimits, isEnterprise, isFree } from '@/lib/billing/plan-helpers'
+import { isOrgScopedSubscription } from '@/lib/billing/subscriptions/utils'
+import { getEnv } from '@/lib/core/config/env'
+import { isBillingEnabled } from '@/lib/core/config/env-flags'
+
+const logger = createLogger('StorageLimits')
+
+/** Resolve the highest-priority subscription via a deferred import (avoids a static cycle). */
+async function resolveSub(userId: string): Promise<HighestPrioritySubscription | null> {
+  const { getHighestPrioritySubscription } = await import('@/lib/billing/core/subscription')
+  return getHighestPrioritySubscription(userId)
+}
+
+/**
+ * Convert GB to bytes
+ */
+function gbToBytes(gb: number): number {
+  return gb * 1024 * 1024 * 1024
+}
+
+/**
+ * Get storage limits from environment variables with fallback to constants
+ * Returns limits in bytes
+ */
+export function getStorageLimits() {
+  return {
+    free: gbToBytes(
+      Number.parseInt(getEnv('FREE_STORAGE_LIMIT_GB') || String(DEFAULT_FREE_STORAGE_LIMIT_GB))
+    ),
+    pro: gbToBytes(
+      Number.parseInt(getEnv('PRO_STORAGE_LIMIT_GB') || String(DEFAULT_PRO_STORAGE_LIMIT_GB))
+    ),
+    team: gbToBytes(
+      Number.parseInt(getEnv('TEAM_STORAGE_LIMIT_GB') || String(DEFAULT_TEAM_STORAGE_LIMIT_GB))
+    ),
+    enterpriseDefault: gbToBytes(
+      Number.parseInt(
+        getEnv('ENTERPRISE_STORAGE_LIMIT_GB') || String(DEFAULT_ENTERPRISE_STORAGE_LIMIT_GB)
+      )
+    ),
+  }
+}
+
+/**
+ * Get storage limit for a specific plan
+ * Returns limit in bytes
+ */
+export function getStorageLimitForPlan(plan: string, metadata?: any): number {
+  const limits = getStorageLimits()
+
+  if (isEnterprise(plan)) {
+    if (metadata?.storageLimitGB) {
+      return gbToBytes(Number.parseInt(metadata.storageLimitGB))
+    }
+    return limits.enterpriseDefault
+  }
+
+  const effectivePlan = getPlanTypeForLimits(plan)
+  const limitByPlan: Record<'free' | 'pro' | 'team', number> = {
+    free: limits.free,
+    pro: limits.pro,
+    team: limits.team,
+  }
+  return limitByPlan[effectivePlan as 'free' | 'pro' | 'team'] ?? limits.free
+}
+
+/**
+ * Get storage limit for a user based on their subscription. Returns limit in
+ * bytes.
+ *
+ * @param prefetchedSub - Pass an already-resolved subscription (may be `null`)
+ *   to skip the `getHighestPrioritySubscription` lookup on hot paths. Omit
+ *   (leave `undefined`) to fetch it here.
+ */
+export async function getUserStorageLimit(
+  userId: string,
+  prefetchedSub?: HighestPrioritySubscription | null
+): Promise<number> {
+  try {
+    const sub = prefetchedSub === undefined ? await resolveSub(userId) : prefetchedSub
+
+    const limits = getStorageLimits()
+
+    if (!sub || isFree(sub.plan)) {
+      return limits.free
+    }
+
+    if (isOrgScopedSubscription(sub, userId)) {
+      const metadata = sub.metadata as { customStorageLimitGB?: number } | null
+      if (metadata?.customStorageLimitGB) {
+        return metadata.customStorageLimitGB * 1024 * 1024 * 1024
+      }
+      return isEnterprise(sub.plan) ? limits.enterpriseDefault : limits.team
+    }
+
+    // Personally-scoped plans use the per-plan default storage cap.
+    const effectivePlan = getPlanTypeForLimits(sub.plan)
+    const limitByPlan: Record<'free' | 'pro' | 'team', number> = {
+      free: limits.free,
+      pro: limits.pro,
+      team: limits.team,
+    }
+    return limitByPlan[effectivePlan as 'free' | 'pro' | 'team'] ?? limits.free
+  } catch (error) {
+    logger.error('Error getting user storage limit:', error)
+    return getStorageLimits().free
+  }
+}
+
+/**
+ * Get current storage usage for a user. Returns usage in bytes.
+ *
+ * @param prefetchedSub - Pass an already-resolved subscription (may be `null`)
+ *   to skip the `getHighestPrioritySubscription` lookup on hot paths.
+ */
+export async function getUserStorageUsage(
+  userId: string,
+  prefetchedSub?: HighestPrioritySubscription | null
+): Promise<number> {
+  try {
+    const sub = prefetchedSub === undefined ? await resolveSub(userId) : prefetchedSub
+
+    // Org-scoped subs share pooled `organization.storageUsedBytes`;
+    // personal plans use `userStats`.
+    if (isOrgScopedSubscription(sub, userId) && sub) {
+      const orgRecord = await db
+        .select({ storageUsedBytes: organization.storageUsedBytes })
+        .from(organization)
+        .where(eq(organization.id, sub.referenceId))
+        .limit(1)
+
+      return orgRecord.length > 0 ? orgRecord[0].storageUsedBytes || 0 : 0
+    }
+
+    const stats = await db
+      .select({ storageUsedBytes: userStats.storageUsedBytes })
+      .from(userStats)
+      .where(eq(userStats.userId, userId))
+      .limit(1)
+
+    return stats.length > 0 ? stats[0].storageUsedBytes || 0 : 0
+  } catch (error) {
+    logger.error('Error getting user storage usage:', error)
+    return 0
+  }
+}
+
+/**
+ * Check if user has storage quota available
+ * Always allows uploads when billing is disabled
+ */
+export async function checkStorageQuota(
+  userId: string,
+  additionalBytes: number
+): Promise<{ allowed: boolean; currentUsage: number; limit: number; error?: string }> {
+  if (!isBillingEnabled) {
+    return {
+      allowed: true,
+      currentUsage: 0,
+      limit: Number.MAX_SAFE_INTEGER,
+    }
+  }
+
+  try {
+    const [currentUsage, limit] = await Promise.all([
+      getUserStorageUsage(userId),
+      getUserStorageLimit(userId),
+    ])
+
+    const newUsage = currentUsage + additionalBytes
+    const allowed = newUsage <= limit
+
+    return {
+      allowed,
+      currentUsage,
+      limit,
+      error: allowed
+        ? undefined
+        : `Storage limit exceeded. Used: ${(newUsage / (1024 * 1024 * 1024)).toFixed(2)}GB, Limit: ${(limit / (1024 * 1024 * 1024)).toFixed(0)}GB`,
+    }
+  } catch (error) {
+    logger.error('Error checking storage quota:', error)
+    return {
+      allowed: false,
+      currentUsage: 0,
+      limit: 0,
+      error: 'Failed to check storage quota',
+    }
+  }
+}

apps/sim/lib/billing/storage/tracking.ts

@@ -0,0 +1,289 @@
+/**
+ * Storage usage tracking
+ * Updates storage_used_bytes for users and organizations
+ * Only tracks when billing is enabled
+ */
+
+import { db } from '@sim/db'
+import { organization, userStats } from '@sim/db/schema'
+import { createLogger } from '@sim/logger'
+import { generateId } from '@sim/utils/id'
+import { and, eq, sql } from 'drizzle-orm'
+import { maybeNotifyLimit } from '@/lib/billing/core/limit-notifications'
+import type { HighestPrioritySubscription } from '@/lib/billing/core/plan'
+import { getUserStorageLimit, getUserStorageUsage } from '@/lib/billing/storage/limits'
+import { getFreeTierLimit, isOrgScopedSubscription } from '@/lib/billing/subscriptions/utils'
+import { isBillingEnabled } from '@/lib/core/config/env-flags'
+
+const logger = createLogger('StorageTracking')
+
+/** Format bytes as a `GB` label for usage-limit emails (2dp usage, whole-number limit). */
+function formatGb(bytes: number, decimals: number): string {
+  return `${(bytes / 1024 ** 3).toFixed(decimals)} GB`
+}
+
+/**
+ * Best-effort storage threshold evaluation after a usage change. Re-reads the
+ * (now updated) usage and plan limit, then delegates dedup + send to
+ * {@link maybeNotifyLimit}. Never throws.
+ *
+ * The caller passes the subscription it already resolved for the increment/
+ * decrement, so the whole path (usage read, limit, scope) reuses a single
+ * `getHighestPrioritySubscription` instead of re-fetching it three times.
+ *
+ * @param rearmOnly - True on decrements, so a shrink that leaves usage above a
+ *   threshold re-arms but never sends (a drop is not a fresh crossing).
+ */
+export async function maybeNotifyStorageLimit(
+  userId: string,
+  workspaceId: string,
+  sub: HighestPrioritySubscription | null,
+  rearmOnly = false
+): Promise<void> {
+  try {
+    const [usage, limit] = await Promise.all([
+      getUserStorageUsage(userId, sub),
+      getUserStorageLimit(userId, sub),
+    ])
+
+    await maybeNotifyLimit({
+      category: 'storage',
+      billedUserId: userId,
+      workspaceId,
+      currentUsage: usage,
+      limit,
+      usageLabel: formatGb(usage, 2),
+      limitLabel: formatGb(limit, 0),
+      rearmOnly,
+      subscription: sub,
+    })
+  } catch (error) {
+    logger.error('Error evaluating storage limit notification:', error)
+  }
+}
+
+/**
+ * Increment storage usage after successful file upload
+ * Only tracks if billing is enabled
+ *
+ * @param workspaceId - When provided, evaluates the storage usage-limit email
+ *   (80% / 100%) after the increment. Best-effort; never blocks the upload.
+ */
+export async function incrementStorageUsage(
+  userId: string,
+  bytes: number,
+  workspaceId?: string
+): Promise<void> {
+  if (!isBillingEnabled) {
+    logger.debug('Billing disabled, skipping storage increment')
+    return
+  }
+
+  let sub: HighestPrioritySubscription | null = null
+  try {
+    const { getHighestPrioritySubscription } = await import('@/lib/billing/core/subscription')
+    sub = await getHighestPrioritySubscription(userId)
+
+    // Org-scoped subs pool at the org level; personal plans per-user.
+    if (isOrgScopedSubscription(sub, userId) && sub) {
+      await db
+        .update(organization)
+        .set({
+          storageUsedBytes: sql`${organization.storageUsedBytes} + ${bytes}`,
+        })
+        .where(eq(organization.id, sub.referenceId))
+
+      logger.info(`Incremented org storage: ${bytes} bytes for org ${sub.referenceId}`)
+    } else {
+      await db
+        .update(userStats)
+        .set({
+          storageUsedBytes: sql`${userStats.storageUsedBytes} + ${bytes}`,
+        })
+        .where(eq(userStats.userId, userId))
+
+      logger.info(`Incremented user storage: ${bytes} bytes for user ${userId}`)
+    }
+  } catch (error) {
+    logger.error('Error incrementing storage usage:', error)
+    throw error
+  }
+
+  if (workspaceId) {
+    void maybeNotifyStorageLimit(userId, workspaceId, sub)
+  }
+}
+
+/**
+ * Decrement storage usage after file deletion
+ * Only tracks if billing is enabled
+ *
+ * @param workspaceId - When provided, re-evaluates the storage threshold state
+ *   after the decrement. Usage only drops here, so this can only re-arm a
+ *   previously-sent threshold (it never sends), keeping the re-warning correct
+ *   after a shrink. Best-effort; never blocks the caller.
+ */
+export async function decrementStorageUsage(
+  userId: string,
+  bytes: number,
+  workspaceId?: string
+): Promise<void> {
+  if (!isBillingEnabled) {
+    logger.debug('Billing disabled, skipping storage decrement')
+    return
+  }
+
+  let sub: HighestPrioritySubscription | null = null
+  try {
+    const { getHighestPrioritySubscription } = await import('@/lib/billing/core/subscription')
+    sub = await getHighestPrioritySubscription(userId)
+
+    if (isOrgScopedSubscription(sub, userId) && sub) {
+      await db
+        .update(organization)
+        .set({
+          storageUsedBytes: sql`GREATEST(0, ${organization.storageUsedBytes} - ${bytes})`,
+        })
+        .where(eq(organization.id, sub.referenceId))
+
+      logger.info(`Decremented org storage: ${bytes} bytes for org ${sub.referenceId}`)
+    } else {
+      await db
+        .update(userStats)
+        .set({
+          storageUsedBytes: sql`GREATEST(0, ${userStats.storageUsedBytes} - ${bytes})`,
+        })
+        .where(eq(userStats.userId, userId))
+
+      logger.info(`Decremented user storage: ${bytes} bytes for user ${userId}`)
+    }
+  } catch (error) {
+    logger.error('Error decrementing storage usage:', error)
+    throw error
+  }
+
+  if (workspaceId) {
+    void maybeNotifyStorageLimit(userId, workspaceId, sub, true)
+  }
+}
+
+type StorageTransaction = Parameters<Parameters<typeof db.transaction>[0]>[0]
+
+/**
+ * Decrement a user's (or their org's) storage counter inside an existing
+ * transaction, using a pre-resolved subscription. This lets a caller make the
+ * counter update atomic with the DB rows it is deleting (e.g. hard-deleting
+ * documents), so a failure of either rolls back both — no inflated counter, no
+ * over-decrement. The caller resolves the subscription (a read) before opening
+ * the transaction.
+ */
+export async function decrementStorageUsageInTx(
+  tx: StorageTransaction,
+  sub: HighestPrioritySubscription | null,
+  userId: string,
+  bytes: number
+): Promise<void> {
+  if (!isBillingEnabled || bytes <= 0) return
+  if (isOrgScopedSubscription(sub, userId) && sub) {
+    await tx
+      .update(organization)
+      .set({ storageUsedBytes: sql`GREATEST(0, ${organization.storageUsedBytes} - ${bytes})` })
+      .where(eq(organization.id, sub.referenceId))
+  } else {
+    await tx
+      .update(userStats)
+      .set({ storageUsedBytes: sql`GREATEST(0, ${userStats.storageUsedBytes} - ${bytes})` })
+      .where(eq(userStats.userId, userId))
+  }
+}
+
+/**
+ * Atomically check quota and increment a user's (or their org's) storage
+ * counter inside an existing transaction, using a pre-resolved subscription.
+ * The check and the increment are a single conditional `UPDATE`, so two
+ * concurrent callers can no longer both read the same pre-increment usage,
+ * both pass the check, and both commit past the limit — the second caller's
+ * `UPDATE` re-evaluates the WHERE clause against the first caller's already
+ * -committed-within-the-same-DB-round-trip row and correctly fails. Replaces
+ * the old read-then-decide-then-increment-after-commit split (`checkStorageQuota`
+ * + a fire-and-forget `incrementStorageUsage` after the transaction), which left
+ * a window between the read and the increment.
+ *
+ * On success, callers should best-effort call {@link maybeNotifyStorageLimit}
+ * after the transaction commits (mirrors the existing post-increment threshold
+ * check) — this helper doesn't do it itself since it runs mid-transaction.
+ *
+ * For a personal (non-org-scoped) `userId`, this first upserts the
+ * `userStats` row on `tx` — a documented possibility for OAuth account
+ * linking (see `ensureUserStatsExists` in `lib/billing/core/usage.ts`, whose
+ * insert values this mirrors) — because the conditional `UPDATE` below
+ * matches 0 rows, and therefore reads as "quota exceeded", if that row
+ * doesn't exist yet. `ensureUserStatsExists` itself isn't reused here since
+ * it writes through the standalone `db` client, which would open a second
+ * pooled connection while this transaction's is held.
+ */
+export async function checkAndIncrementStorageUsageInTx(
+  tx: StorageTransaction,
+  sub: HighestPrioritySubscription | null,
+  userId: string,
+  bytes: number
+): Promise<{ allowed: boolean; currentUsage: number; limit: number; error?: string }> {
+  if (!isBillingEnabled) {
+    return { allowed: true, currentUsage: 0, limit: Number.MAX_SAFE_INTEGER }
+  }
+
+  const limit = await getUserStorageLimit(userId, sub)
+
+  if (bytes <= 0) {
+    return { allowed: true, currentUsage: await getUserStorageUsage(userId, sub), limit }
+  }
+
+  const orgScoped = isOrgScopedSubscription(sub, userId) && sub
+
+  if (!orgScoped) {
+    await tx
+      .insert(userStats)
+      .values({
+        id: generateId(),
+        userId,
+        currentUsageLimit: getFreeTierLimit().toString(),
+        usageLimitUpdatedAt: new Date(),
+      })
+      .onConflictDoNothing({ target: userStats.userId })
+  }
+
+  const [updated] = orgScoped
+    ? await tx
+        .update(organization)
+        .set({ storageUsedBytes: sql`${organization.storageUsedBytes} + ${bytes}` })
+        .where(
+          and(
+            eq(organization.id, sub.referenceId),
+            sql`${organization.storageUsedBytes} + ${bytes} <= ${limit}`
+          )
+        )
+        .returning({ storageUsedBytes: organization.storageUsedBytes })
+    : await tx
+        .update(userStats)
+        .set({ storageUsedBytes: sql`${userStats.storageUsedBytes} + ${bytes}` })
+        .where(
+          and(
+            eq(userStats.userId, userId),
+            sql`${userStats.storageUsedBytes} + ${bytes} <= ${limit}`
+          )
+        )
+        .returning({ storageUsedBytes: userStats.storageUsedBytes })
+
+  if (updated) {
+    return { allowed: true, currentUsage: updated.storageUsedBytes - bytes, limit }
+  }
+
+  const currentUsage = await getUserStorageUsage(userId, sub)
+  const newUsage = currentUsage + bytes
+  return {
+    allowed: false,
+    currentUsage,
+    limit,
+    error: `Storage limit exceeded. Used: ${(newUsage / (1024 * 1024 * 1024)).toFixed(2)}GB, Limit: ${(limit / (1024 * 1024 * 1024)).toFixed(0)}GB`,
+  }
+}