improvement(billing): key the usage email claim on period and limit, and harden the shared read

Author: waleedlatif1

apps/sim/lib/billing/core/limit-notifications.ts

@@ -2,7 +2,7 @@ import { db } from '@sim/db'
 import { member, organization, settings, user, userStats } from '@sim/db/schema'
 import { createLogger } from '@sim/logger'
 import { isOrgAdminRole } from '@sim/platform-authz/workspace'
-import { and, eq, sql } from 'drizzle-orm'
+import { and, eq, type SQL, sql } from 'drizzle-orm'
 import type { HighestPrioritySubscription } from '@/lib/billing/core/plan'
 import { getHighestPrioritySubscription } from '@/lib/billing/core/subscription'
 import type { BillingEntity } from '@/lib/billing/core/usage-log'
@@ -17,9 +17,6 @@ const logger = createLogger('LimitNotifications')
 /** Limit categories that send per-category threshold emails (credits has its own path). */
 export type LimitCategory = Extract<UpgradeReason, 'storage' | 'tables' | 'seats'>
 
-/** Every category whose emailed threshold is persisted, including credits. */
-type ClaimCategory = LimitCategory | Extract<UpgradeReason, 'credits'>
-
 const WARN_THRESHOLD = 80
 const REACH_THRESHOLD = 100
 /** Usage must drop below this band before the same threshold can re-notify (hysteresis). */
@@ -35,6 +32,32 @@ function thresholdFor(percent: number): 0 | 80 | 100 {
   return 0
 }
 
+/**
+ * Replace the account's `limitNotifications` with `next` when `condition` holds, returning
+ * whether the row was updated.
+ */
+async function writeLimitNotifications(
+  scope: 'user' | 'organization',
+  id: string,
+  next: SQL,
+  condition: SQL
+): Promise<boolean> {
+  const written =
+    scope === 'user'
+      ? await db
+          .update(userStats)
+          .set({ limitNotifications: next })
+          .where(and(eq(userStats.userId, id), condition))
+          .returning({ id: userStats.userId })
+      : await db
+          .update(organization)
+          .set({ limitNotifications: next })
+          .where(and(eq(organization.id, id), condition))
+          .returning({ id: organization.id })
+
+  return written.length > 0
+}
+
 /**
  * Atomically claim a threshold for a category: advance the stored value to
  * `threshold` only if it is currently lower, returning whether THIS call won the
@@ -44,7 +67,7 @@ function thresholdFor(percent: number): 0 | 80 | 100 {
 async function claimThreshold(
   scope: 'user' | 'organization',
   id: string,
-  category: ClaimCategory,
+  category: LimitCategory,
   threshold: number
 ): Promise<boolean> {
   const setExpr = sql`jsonb_set(coalesce(${scope === 'user' ? userStats.limitNotifications : organization.limitNotifications}, '{}'::jsonb), ARRAY[${category}], to_jsonb(${threshold}::int))`
@@ -53,38 +76,37 @@ async function claimThreshold(
       ? sql`coalesce((${userStats.limitNotifications} ->> ${category})::int, 0) < ${threshold}`
       : sql`coalesce((${organization.limitNotifications} ->> ${category})::int, 0) < ${threshold}`
 
-  const claimed =
-    scope === 'user'
-      ? await db
-          .update(userStats)
-          .set({ limitNotifications: setExpr })
-          .where(and(eq(userStats.userId, id), onlyIfLower))
-          .returning({ id: userStats.userId })
-      : await db
-          .update(organization)
-          .set({ limitNotifications: setExpr })
-          .where(and(eq(organization.id, id), onlyIfLower))
-          .returning({ id: organization.id })
-
-  return claimed.length > 0
+  return writeLimitNotifications(scope, id, setExpr, onlyIfLower)
 }
 
 const DAY_MS = 24 * 60 * 60 * 1000
 
 /**
- * Claim a credits threshold (80 or 100) once per billing period, returning whether THIS call won
- * it. The stored value is the period's start day followed by the threshold, so it only grows: a
- * later period outranks every claim of an earlier one and re-arms both thresholds with no reset
- * write, while within a period a claim of 100 also retires 80, and never the reverse.
+ * Claim a credits threshold (80 or 100), returning whether THIS call won it. The claim is keyed
+ * on the billing period and the limit: `credits` holds the highest threshold emailed while
+ * `creditsPeriod` (the period's start day) and `creditsLimit` (the limit in cents) still match,
+ * so a new period or a changed limit — in either direction — re-arms both thresholds with no
+ * reset write, while within one a claim of 100 also retires 80, and never the reverse.
  */
-export function claimCreditsThreshold(
-  scope: 'user' | 'organization',
-  id: string,
-  periodStart: Date,
+export function claimCreditsThreshold(params: {
+  scope: 'user' | 'organization'
+  id: string
+  periodStart: Date
+  limit: number
   threshold: 80 | 100
-): Promise<boolean> {
-  const periodDay = Math.floor(periodStart.getTime() / DAY_MS)
-  return claimThreshold(scope, id, 'credits', periodDay * 1000 + threshold)
+}): Promise<boolean> {
+  const { scope, id, threshold } = params
+  const periodDay = Math.floor(params.periodStart.getTime() / DAY_MS)
+  const limitCents = Math.round(params.limit * 100)
+  const column = scope === 'user' ? userStats.limitNotifications : organization.limitNotifications
+  const next = sql`coalesce(${column}, '{}'::jsonb) || jsonb_build_object('credits', ${threshold}::int, 'creditsPeriod', ${periodDay}::bigint, 'creditsLimit', ${limitCents}::bigint)`
+  const unclaimed = sql`not (
+    (${column} ->> 'creditsPeriod')::bigint is not distinct from ${periodDay}::bigint
+    and (${column} ->> 'creditsLimit')::bigint is not distinct from ${limitCents}::bigint
+    and coalesce((${column} ->> 'credits')::int, 0) >= ${threshold}::int
+  )`
+
+  return writeLimitNotifications(scope, id, next, unclaimed)
 }
 
 /** Re-arm a category (reset its stored threshold to 0) once usage falls back into the low band. */
@@ -129,7 +151,7 @@ async function isUnsubscribed(email: string): Promise<boolean> {
  * Returning an empty list means "nobody to notify" — the caller then skips the
  * claim so the dedup state isn't burned without an email going out.
  */
-async function resolveRecipients(
+export async function resolveLimitEmailRecipients(
   scope: 'user' | 'organization',
   params: { userId?: string; userEmail?: string; userName?: string; organizationId?: string }
 ): Promise<LimitEmailRecipient[]> {
@@ -228,7 +250,7 @@ export async function maybeSendLimitThresholdEmail(params: {
 
     if (params.rearmOnly || desired === 0) return
 
-    const recipients = await resolveRecipients(scope, params)
+    const recipients = await resolveLimitEmailRecipients(scope, params)
     if (recipients.length === 0) return
 
     if (!(await claimThreshold(scope, stateId, category, desired))) return

apps/sim/lib/billing/core/reporting-usage-cache.integration.ts

@@ -95,11 +95,30 @@ describe.runIf(Boolean(redisUrl))('shared reporting usage read', () => {
     expect(ttl).toBeLessThanOrEqual(35_000)
   })
 
-  it('treats an unreadable stored sum as a miss and overwrites it', async () => {
+  it('treats an unreadable stored sum as a miss', async () => {
     await redis.set(sharedKey(payer), 'not-a-number', 'PX', 30_000)
 
     await expect(readSoftGateUsageCost(payer, REPORTING)).resolves.toBe(5.75)
-    await vi.waitFor(async () => expect(await redis.get(sharedKey(payer))).toBe('5.75'))
+    expect(transaction).toHaveBeenCalledTimes(1)
+  })
+
+  it('never lets a slower, older sum replace one stored while it ran', async () => {
+    transaction.mockImplementationOnce(async (callback) => {
+      const result = await database.transaction(callback)
+      await redis.set(sharedKey(payer), '9.99', 'PX', 30_000)
+      return result
+    })
+
+    await expect(readSoftGateUsageCost(payer, REPORTING)).resolves.toBe(5.75)
+    expect(await redis.get(sharedKey(payer))).toBe('9.99')
+  })
+
+  it('sums the ledger when the Redis client cannot be built', async () => {
+    redisConfigMockFns.mockGetRedisClient.mockImplementation(() => {
+      throw new Error('Invalid Redis configuration')
+    })
+
+    await expect(readSoftGateUsageCost(payer, REPORTING)).resolves.toBe(5.75)
   })
 
   it('sums the ledger promptly when Redis is unreachable', async () => {

apps/sim/lib/billing/core/reporting-usage-cache.ts

@@ -26,8 +26,10 @@ const logger = createLogger('ReportingUsageCache')
  * above the limit is a refusal the true sum would also give. Thirty seconds keeps that overrun
  * small against a year-long allowance while turning a per-event scan into one per window.
  *
- * A sum is held both in Redis, shared by every process, and in each process that reads it, so a
- * served sum can be up to twice this old (plus the Redis expiry's jitter).
+ * A sum is held both in Redis, shared by every process, and in each process that reads it. It
+ * reflects the ledger as of the moment its sum began, so a served sum can omit usage written over
+ * the sum's own duration, plus up to this TTL and its jitter in Redis, plus up to this TTL again
+ * in the reading process.
  */
 export const REPORTING_USAGE_CACHE_TTL_MS = 30_000
 
@@ -57,9 +59,9 @@ function sharedReportingUsageKey(key: string): string {
  * sums the ledger, so the cache can cost a read its latency but never its answer.
  */
 async function readSharedReportingUsageCost(key: string): Promise<number | undefined> {
-  const redis = getRedisClient()
-  if (!redis) return undefined
   try {
+    const redis = getRedisClient()
+    if (!redis) return undefined
     const stored = await withinDeadline(
       () => redis.get(sharedReportingUsageKey(key)),
       Date.now() + SHARED_READ_TIMEOUT_MS
@@ -76,14 +78,26 @@ async function readSharedReportingUsageCost(key: string): Promise<number | undef
   return undefined
 }
 
-/** Fire-and-forget: a read never waits on, or fails because of, the shared write. */
+function warnSharedWriteFailed(error: unknown): void {
+  logger.warn('Shared reporting usage write failed', { error: getErrorMessage(error) })
+}
+
+/**
+ * Fire-and-forget: a read never waits on, or fails because of, the shared write. The write only
+ * lands when no sum is stored (`NX`), so a slow, older sum can never replace a fresher one or
+ * extend its expiry.
+ */
 function writeSharedReportingUsageCost(key: string, cost: number): void {
-  const redis = getRedisClient()
-  if (!redis) return
-  const ttlMs = REPORTING_USAGE_CACHE_TTL_MS + randomInt(0, SHARED_TTL_JITTER_MS)
-  redis.set(sharedReportingUsageKey(key), String(cost), 'PX', ttlMs).catch((error: unknown) => {
-    logger.warn('Shared reporting usage write failed', { error: getErrorMessage(error) })
-  })
+  try {
+    const redis = getRedisClient()
+    if (!redis) return
+    const ttlMs = REPORTING_USAGE_CACHE_TTL_MS + randomInt(0, SHARED_TTL_JITTER_MS)
+    redis
+      .set(sharedReportingUsageKey(key), String(cost), 'PX', ttlMs, 'NX')
+      .catch(warnSharedWriteFailed)
+  } catch (error) {
+    warnSharedWriteFailed(error)
+  }
 }
 
 /**
@@ -153,9 +167,10 @@ async function readCachedReportingUsageCost(
 /**
  * Period usage for a soft reader: an admission check, a display, or a level-triggered
  * notification that tolerates the cache's bounded under-count. Enterprise reporting windows are
- * served from the shared cache for up to twice {@link REPORTING_USAGE_CACHE_TTL_MS}, since their
- * year-long sum is the expensive one; every other period is summed exactly, as before. A read on
- * a caller's own executor (a transaction or a replica) keeps its own snapshot and is never shared.
+ * served from the shared cache, within the lag {@link REPORTING_USAGE_CACHE_TTL_MS} describes,
+ * since their year-long sum is the expensive one; every other period is summed exactly, as
+ * before. A read on a caller's own executor (a transaction or a replica) keeps its own snapshot
+ * and is never shared.
  * Never use it for invoicing, cycle close, an edge-triggered decision, or a read that must see its
  * own write.
  */

apps/sim/lib/billing/core/usage-threshold-email.integration.ts

@@ -40,27 +40,37 @@ const SEPTEMBER = new Date('2026-09-01T00:00:00.000Z')
 const OCTOBER = new Date('2026-10-01T00:00:00.000Z')
 
 let organizationId: string
+let adminId: string
 
-function notify(currentUsage: number, periodStart = SEPTEMBER) {
+/** One completion that leaves the organization at `usage`, having recorded `costDelta` of it. */
+function notify(
+  usage: number,
+  { periodStart = SEPTEMBER, limit = 100, costDelta = 1 } = {}
+): Promise<void> {
   return maybeSendUsageThresholdEmail({
     scope: 'organization',
     organizationId,
     planName: 'Enterprise',
     periodStart,
     workspaceId: 'workspace',
-    currentUsage,
-    limit: 100,
+    usageBefore: usage - costDelta,
+    costDelta,
+    limit,
   })
 }
 
+async function setNotificationsEnabled(enabled: boolean): Promise<void> {
+  await connection`INSERT INTO settings VALUES (${adminId}, ${adminId}, ${enabled})
+    ON CONFLICT (id) DO UPDATE SET billing_usage_notifications_enabled = ${enabled}`
+}
+
 beforeAll(async () => {
   await connection.unsafe(`CREATE SCHEMA "${schemaName}"`)
   await connection.unsafe(`
     CREATE TABLE organization (id text PRIMARY KEY, limit_notifications jsonb);
     CREATE TABLE "user" (id text PRIMARY KEY, email text, name text);
     CREATE TABLE member (id text PRIMARY KEY, organization_id text, user_id text, role text);
     CREATE TABLE settings (id text PRIMARY KEY, user_id text, billing_usage_notifications_enabled boolean);
-    INSERT INTO "user" VALUES ('admin', 'admin@example.com', 'Admin');
   `)
   select.mockImplementation((fields) => database.select(fields))
   update.mockImplementation((table) => database.update(table))
@@ -70,8 +80,10 @@ beforeAll(async () => {
 beforeEach(async () => {
   mockSendEmail.mockClear()
   organizationId = generateId()
+  adminId = generateId()
   await connection`INSERT INTO organization (id) VALUES (${organizationId})`
-  await connection`INSERT INTO member VALUES (${generateId()}, ${organizationId}, 'admin', 'owner')`
+  await connection`INSERT INTO "user" VALUES (${adminId}, ${`${adminId}@example.com`}, 'Admin')`
+  await connection`INSERT INTO member VALUES (${generateId()}, ${organizationId}, ${adminId}, 'owner')`
 })
 
 afterAll(async () => {
@@ -97,12 +109,37 @@ describe('usage threshold email', () => {
     expect(mockSendEmail).toHaveBeenCalledTimes(2)
   })
 
-  it('re-arms both thresholds in the next billing period', async () => {
+  it('re-arms both thresholds whenever the billing period changes, even to an earlier one', async () => {
     await notify(100)
-    await notify(85, OCTOBER)
-    await notify(100, OCTOBER)
+    await notify(85, { periodStart: OCTOBER })
+    await notify(100, { periodStart: OCTOBER })
     await notify(85)
 
-    expect(mockSendEmail).toHaveBeenCalledTimes(3)
+    expect(mockSendEmail).toHaveBeenCalledTimes(4)
+  })
+
+  it('warns again at a raised limit after the old one was reached', async () => {
+    await notify(100)
+    await notify(100, { limit: 125 })
+    await notify(110, { limit: 125 })
+
+    expect(mockSendEmail).toHaveBeenCalledTimes(2)
+  })
+
+  it('keeps the claim for a later completion when nobody can be notified', async () => {
+    await setNotificationsEnabled(false)
+    await notify(90)
+    await setNotificationsEnabled(true)
+    await notify(90)
+
+    expect(mockSendEmail).toHaveBeenCalledTimes(1)
+  })
+
+  it('keeps the claim when a completion recorded no cost', async () => {
+    await notify(90, { costDelta: 0 })
+    expect(mockSendEmail).not.toHaveBeenCalled()
+
+    await notify(90)
+    expect(mockSendEmail).toHaveBeenCalledTimes(1)
   })
 })

apps/sim/lib/billing/core/usage.test.ts

@@ -334,7 +334,8 @@ describe('maybeSendUsageThresholdEmail', () => {
   it('emails a paid personal account at 100% with the raise-your-limit template', async () => {
     await maybeSendUsageThresholdEmail({
       ...paidUser,
-      currentUsage: 20,
+      usageBefore: 19,
+      costDelta: 1,
     })
 
     expect(mockRenderUsageLimitReached).toHaveBeenCalledWith(
@@ -359,7 +360,8 @@ describe('maybeSendUsageThresholdEmail', () => {
       organizationId: 'org-1',
       workspaceId: 'ws-1',
       periodStart: new Date('2026-09-01T00:00:00.000Z'),
-      currentUsage: 500,
+      usageBefore: 499,
+      costDelta: 1,
       limit: 500,
     })