fix(billing): re-arm limit-notification dedup on usage drops (prior-usage + decrement)

waleedlatif1 · waleedlatif1 · commit fb99d84dd450 · 2026-06-22T13:59:10.000-07:00
diff --git a/apps/sim/lib/billing/core/limit-notifications.test.ts b/apps/sim/lib/billing/core/limit-notifications.test.ts
@@ -7,6 +7,7 @@ const {
   billingFlag,
   mockClaim,
   mockSelectRows,
+  dbUpdateSpy,
   sendEmailSpy,
   getEmailPreferencesMock,
   renderMock,
@@ -16,6 +17,7 @@ const {
   billingFlag: { enabled: true },
   mockClaim: vi.fn<[], unknown[]>(() => [{ id: 'u1' }]),
   mockSelectRows: vi.fn<[], unknown[]>(() => []),
+  dbUpdateSpy: vi.fn(),
   sendEmailSpy: vi.fn(() => Promise.resolve({ success: true })),
   getEmailPreferencesMock: vi.fn(() => Promise.resolve(null as unknown)),
   renderMock: vi.fn(() => Promise.resolve('<html></html>')),
@@ -42,7 +44,8 @@ vi.mock('@sim/db', () => {
     then: (f: (v: unknown) => unknown, r?: (e: unknown) => unknown) =>
       Promise.resolve(mockSelectRows()).then(f, r),
   }
-  return { db: { update: () => updateBuilder, select: () => selectBuilder } }
+  dbUpdateSpy.mockImplementation(() => updateBuilder)
+  return { db: { update: dbUpdateSpy, select: () => selectBuilder } }
 })
 
 vi.mock('@/lib/core/config/env-flags', () => ({
@@ -102,6 +105,31 @@ describe('maybeSendLimitThresholdEmail', () => {
     expect(sendEmailSpy).not.toHaveBeenCalled()
   })
 
+  it('re-arms then claims when a single jump crosses from below the band past 80% (priorUsage)', async () => {
+    // prior 50% (re-arm band) → current 90%: re-arm (update) + claim (update) = 2 updates, then send.
+    await maybeSendLimitThresholdEmail({
+      ...baseUserParams,
+      currentUsage: 4.5,
+      limit: 5,
+      priorUsage: 2.5,
+    })
+    expect(dbUpdateSpy).toHaveBeenCalledTimes(2)
+    expect(sendEmailSpy).toHaveBeenCalledTimes(1)
+    expect(renderMock).toHaveBeenCalledWith(expect.objectContaining({ kind: 'warning' }))
+  })
+
+  it('does not re-arm when prior usage was already in-band (no spurious reset)', async () => {
+    // prior 85% and current 90%: both >= re-arm band → claim only, no re-arm update.
+    await maybeSendLimitThresholdEmail({
+      ...baseUserParams,
+      currentUsage: 4.5,
+      limit: 5,
+      priorUsage: 4.25,
+    })
+    expect(dbUpdateSpy).toHaveBeenCalledTimes(1)
+    expect(sendEmailSpy).toHaveBeenCalledTimes(1)
+  })
+
   it('does not send in the dead band (70%–80%)', async () => {
     await maybeSendLimitThresholdEmail({ ...baseUserParams, currentUsage: 3.75, limit: 5 })
     expect(mockClaim).not.toHaveBeenCalled()
diff --git a/apps/sim/lib/billing/core/limit-notifications.ts b/apps/sim/lib/billing/core/limit-notifications.ts
@@ -115,6 +115,13 @@ export async function maybeSendLimitThresholdEmail(params: {
   usageLabel: string
   /** Pre-formatted limit for the email body, e.g. "5 GB", "10 seats". */
   limitLabel: string
+  /**
+   * Usage immediately BEFORE the mutation, when known (e.g. the pre-insert row
+   * count). Lets a single large change that jumps from below the re-arm band
+   * past a threshold still re-arm before claiming, so the re-warning isn't
+   * suppressed. Omit when only the post-mutation usage is observable.
+   */
+  priorUsage?: number
   userId?: string
   userEmail?: string
   userName?: string
@@ -126,14 +133,18 @@ export async function maybeSendLimitThresholdEmail(params: {
 
     const { category, scope } = params
     const percent = (params.currentUsage / params.limit) * 100
+    const priorPercent =
+      params.priorUsage != null && params.priorUsage > 0
+        ? (params.priorUsage / params.limit) * 100
+        : percent
     const desired = thresholdFor(percent)
 
     const stateId = scope === 'user' ? params.userId : params.organizationId
     if (!stateId) return
 
-    if (percent < REARM_BELOW) {
+    // Re-arm if usage is (or just was) back in the low band, so a fresh climb re-notifies.
+    if (Math.min(percent, priorPercent) < REARM_BELOW) {
       await rearmThreshold(scope, stateId, category)
-      return
     }
 
     if (desired === 0) return
@@ -229,6 +240,8 @@ export async function maybeNotifyLimit(params: {
   limit: number
   usageLabel: string
   limitLabel: string
+  /** Usage before the mutation, when known — see {@link maybeSendLimitThresholdEmail}. */
+  priorUsage?: number
 }): Promise<void> {
   try {
     const sub = await getHighestPrioritySubscription(params.billedUserId)
@@ -254,6 +267,7 @@ export async function maybeNotifyLimit(params: {
       limit: params.limit,
       usageLabel: params.usageLabel,
       limitLabel: params.limitLabel,
+      priorUsage: params.priorUsage,
       userId: params.billedUserId,
       userEmail,
       userName,
diff --git a/apps/sim/lib/billing/storage/tracking.ts b/apps/sim/lib/billing/storage/tracking.ts
@@ -101,8 +101,17 @@ export async function incrementStorageUsage(
 /**
  * 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): Promise<void> {
+export async function decrementStorageUsage(
+  userId: string,
+  bytes: number,
+  workspaceId?: string
+): Promise<void> {
   if (!isBillingEnabled) {
     logger.debug('Billing disabled, skipping storage decrement')
     return
@@ -135,4 +144,8 @@ export async function decrementStorageUsage(userId: string, bytes: number): Prom
     logger.error('Error decrementing storage usage:', error)
     throw error
   }
+
+  if (workspaceId) {
+    void maybeNotifyStorageLimit(userId, workspaceId)
+  }
 }
diff --git a/apps/sim/lib/table/billing.ts b/apps/sim/lib/table/billing.ts
@@ -23,6 +23,7 @@ const TABLE_ROW_NOTIFY_PERCENT = 80
  */
 async function maybeNotifyTableRowLimit(
   workspaceId: string,
+  currentRowCount: number,
   projectedRowCount: number,
   limit: number
 ): Promise<void> {
@@ -38,6 +39,7 @@ async function maybeNotifyTableRowLimit(
       limit,
       usageLabel: `${projectedRowCount.toLocaleString('en-US')} rows`,
       limitLabel: `${limit.toLocaleString('en-US')} rows`,
+      priorUsage: currentRowCount,
     })
   } catch (error) {
     logger.error('Error evaluating table row-limit notification:', error)
@@ -180,7 +182,7 @@ export async function assertRowCapacity(params: {
   if (limit > 0) {
     const projected = params.currentRowCount + params.addedRows
     if ((projected / limit) * 100 >= TABLE_ROW_NOTIFY_PERCENT) {
-      void maybeNotifyTableRowLimit(params.workspaceId, projected, limit)
+      void maybeNotifyTableRowLimit(params.workspaceId, params.currentRowCount, projected, limit)
     }
   }
 }
diff --git a/apps/sim/lib/uploads/contexts/workspace/workspace-file-manager.ts b/apps/sim/lib/uploads/contexts/workspace/workspace-file-manager.ts
@@ -937,7 +937,7 @@ export async function updateWorkspaceFileContent(
         if (sizeDiff > 0) {
           await incrementStorageUsage(userId, sizeDiff, workspaceId)
         } else {
-          await decrementStorageUsage(userId, Math.abs(sizeDiff))
+          await decrementStorageUsage(userId, Math.abs(sizeDiff), workspaceId)
         }
       } catch (storageError) {
         logger.error(`Failed to update storage tracking:`, storageError)

Original file line number	Diff line number	Diff line change
`@@ -937,7 +937,7 @@ export async function updateWorkspaceFileContent(`
`937`	`937`	`if (sizeDiff > 0) {`
`938`	`938`	`await incrementStorageUsage(userId, sizeDiff, workspaceId)`
`939`	`939`	`} else {`
`940`		`- await decrementStorageUsage(userId, Math.abs(sizeDiff))`
	`940`	`+ await decrementStorageUsage(userId, Math.abs(sizeDiff), workspaceId)`
`941`	`941`	`}`
`942`	`942`	`} catch (storageError) {`
`943`	`943`	logger.error(`Failed to update storage tracking:`, storageError)