fix(storage): decrement KB document storage atomically with deletion

Address Cursor review: hardDeleteDocuments deleted the rows then decremented
best-effort, so a decrement failure left billed storage inflated with no row to
reconcile. Resolve each owner's subscription up front, then decrement inside the
same transaction that deletes the embeddings/documents (decrementStorageUsageInTx,
also now shared by releaseDeletedFileStorage), so the counter and the content
commit or roll back together. Connector docs remain excluded.

Author: waleedlatif1

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

@@ -1,6 +1,7 @@
 export { checkStorageQuota, getUserStorageLimit, getUserStorageUsage } from './limits'
 export {
   decrementStorageUsage,
+  decrementStorageUsageInTx,
   incrementStorageUsage,
   releaseDeletedFileStorage,
 } from './tracking'

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

@@ -166,6 +166,36 @@ export async function decrementStorageUsage(
   }
 }
 
+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 soft-delete a file's metadata row and decrement the owner's storage
  * counter in a single transaction.
@@ -193,7 +223,6 @@ export async function releaseDeletedFileStorage(
 
   const { getHighestPrioritySubscription } = await import('@/lib/billing/core/subscription')
   const sub = await getHighestPrioritySubscription(userId)
-  const orgScoped = isOrgScopedSubscription(sub, userId) && sub !== null
 
   let claimed = false
   await db.transaction(async (tx) => {
@@ -204,18 +233,7 @@ export async function releaseDeletedFileStorage(
       .returning({ id: workspaceFiles.id })
     if (claimedRows.length === 0) return
     claimed = true
-
-    if (orgScoped && 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))
-    }
+    await decrementStorageUsageInTx(tx, sub, userId, bytes)
   })
 
   if (claimed && workspaceId) {

apps/sim/lib/knowledge/documents/service.ts

@@ -14,10 +14,12 @@ import { generateId } from '@sim/utils/id'
 import { tasks } from '@trigger.dev/sdk'
 import { and, asc, desc, eq, inArray, isNotNull, isNull, type SQL, sql } from 'drizzle-orm'
 import { checkActorUsageLimits } from '@/lib/billing/calculations/usage-monitor'
+import type { HighestPrioritySubscription } from '@/lib/billing/core/plan'
+import { getHighestPrioritySubscription } from '@/lib/billing/core/subscription'
 import { recordUsage } from '@/lib/billing/core/usage-log'
 import {
   checkStorageQuota,
-  decrementStorageUsage,
+  decrementStorageUsageInTx,
   incrementStorageUsage,
 } from '@/lib/billing/storage'
 import { checkAndBillOverageThreshold } from '@/lib/billing/threshold-billing'
@@ -2051,28 +2053,36 @@ export async function hardDeleteDocuments(
 
   const existingIds = documentsToDelete.map((doc) => doc.id)
 
+  // Sum the billable bytes per owning user. Connector-synced documents are never
+  // metered at ingest (the sync engine inserts them directly), so they must not
+  // be decremented here.
+  const bytesByUser = new Map<string, number>()
+  for (const doc of documentsToDelete) {
+    if (doc.connectorId || doc.fileSize <= 0) continue
+    const billedUserId = doc.uploadedBy ?? doc.kbUserId
+    if (!billedUserId) continue
+    bytesByUser.set(billedUserId, (bytesByUser.get(billedUserId) ?? 0) + doc.fileSize)
+  }
+
+  // Resolve each owner's subscription before the transaction (these are reads),
+  // then decrement the counters inside the same transaction that deletes the
+  // rows, so a failure of either rolls both back — the billed storage can never
+  // be left inflated once the content is gone.
+  const subByUser = new Map<string, HighestPrioritySubscription | null>()
+  for (const billedUserId of bytesByUser.keys()) {
+    subByUser.set(billedUserId, await getHighestPrioritySubscription(billedUserId))
+  }
+
   await db.transaction(async (tx) => {
     await tx.delete(embedding).where(inArray(embedding.documentId, existingIds))
     await tx.delete(document).where(inArray(document.id, existingIds))
+    for (const [billedUserId, bytes] of bytesByUser) {
+      await decrementStorageUsageInTx(tx, subByUser.get(billedUserId) ?? null, billedUserId, bytes)
+    }
   })
 
   await deleteDocumentStorageFiles(documentsToDelete, requestId)
 
-  for (const doc of documentsToDelete) {
-    // Connector-synced documents are never metered at ingest (the sync engine
-    // inserts them directly), so they must not be decremented here — doing so
-    // would erode legitimately-counted bytes in the same owner counter.
-    if (doc.connectorId) continue
-    const billedUserId = doc.uploadedBy ?? doc.kbUserId
-    if (billedUserId && doc.fileSize > 0) {
-      try {
-        await decrementStorageUsage(billedUserId, doc.fileSize, doc.workspaceId ?? undefined)
-      } catch (storageError) {
-        logger.error(`[${requestId}] Failed to update storage tracking:`, storageError)
-      }
-    }
-  }
-
   logger.info(`[${requestId}] Hard deleted ${existingIds.length} documents`, {
     documentIds: existingIds,
   })