fix(billing): re-read the usage gate often once a payer nears its limit (#9091)

* fix(billing): re-read the usage gate often once a payer nears its limit

The usage gate cached a within-limit admission for five minutes, so a payer
who crossed the limit kept running executions and copilot steps for up to
that long. On a free plan that is unbilled spend; on a capped plan it bills
past the cap. An admission is now held for 10s when the payer or the acting
member has less headroom than the larger of $25 and 20% of the limit (every
free-plan payer), and for the full five minutes otherwise. The direct-v1
account verdict cache follows the same rule.

* test(billing): drive the gate-cache clock with faked timers instead of real waits

Author: waleedlatif1

apps/sim/app/api/billing/update-cost/route.ts

@@ -82,8 +82,8 @@ function invalidBillingProtocolResponse(requestId: string, span: Span): NextResp
  * once it crosses the limit instead of at its next admission, with the card the worker writes
  * to its log. The payer is the attributed run's, or the one a direct-v1 run was admitted with.
  * A duplicate callback answers too: it is often a retry whose first answer was lost. An
- * admission is cached per payer and actor for the gate TTL and a refusal is always re-read, so
- * steady-state steps cost no ledger read. The charge is
+ * admission is cached per payer and actor (briefly once near the limit) and a refusal is always
+ * re-read, so most steps cost no ledger read. The charge is
  * already recorded when this runs; a gate that cannot answer reports not-exceeded and leaves the
  * refusal to the next step or re-check rather than ending a paying run on a database blip,
  * and so does a verdict read that outlasts {@link USAGE_STANDING_TIMEOUT_MS}. An exceeded

apps/sim/app/api/copilot/api-keys/validate/route.test.ts

@@ -6,6 +6,7 @@ import {
   schemaMock,
   setEnvFlags,
 } from '@sim/testing'
+import { createRouteContext } from '@sim/testing/helpers/http'
 import { billingCoreMock, billingCoreMockFns } from '@sim/testing/mocks/billing-core.mock'
 import { billingPlanMock, billingPlanMockFns } from '@sim/testing/mocks/billing-plan.mock'
 import {
@@ -145,7 +146,10 @@ import {
   validateCopilotApiKeyContract,
 } from '@/lib/api/contracts/copilot'
 import { resetMidRunUsageCaches } from '@/lib/billing/core/mid-run-usage'
-import { resetUsageGateCache } from '@/lib/billing/core/usage-gate-cache'
+import {
+  NEAR_LIMIT_USAGE_GATE_TTL_MS,
+  resetUsageGateCache,
+} from '@/lib/billing/core/usage-gate-cache'
 import { POST } from '@/app/api/copilot/api-keys/validate/route'
 
 const { mockGetWorkspaceBillingSettings } = workspacesUtilsMockFns
@@ -629,6 +633,25 @@ describe('validation lifecycle purposes', () => {
     })
   })
 
+  it('re-reads a direct-v1 account near its limit, so a continuation past it is refused', async () => {
+    vi.useFakeTimers({ toFake: ['setTimeout', 'clearTimeout'] })
+    try {
+      mockCheckUsageStatus.mockResolvedValue({ isExceeded: false, currentUsage: 4.2, limit: 5 })
+      expect((await POST(request(body, directHeaders), createRouteContext({}))).status).toBe(200)
+
+      mockCheckUsageStatus.mockResolvedValue({ isExceeded: true, currentUsage: 5.4, limit: 5 })
+      /** `lru-cache` captures `performance` at load: spy its clock, then fire its 1 ms debounce. */
+      const start = performance.now()
+      vi.spyOn(performance, 'now').mockReturnValue(start + NEAR_LIMIT_USAGE_GATE_TTL_MS + 1)
+      await vi.advanceTimersByTimeAsync(1)
+
+      expect((await POST(request(body, directHeaders), createRouteContext({}))).status).toBe(402)
+    } finally {
+      vi.runOnlyPendingTimers()
+      vi.useRealTimers()
+    }
+  })
+
   it('admits a direct-v1 continuation whose usage cannot be read', async () => {
     mockCheckUsageStatus.mockResolvedValue({ isExceeded: true, unavailable: true })
     expect((await POST(request(body, directHeaders))).status).toBe(200)

apps/sim/lib/billing/core/mid-run-usage.ts

@@ -17,6 +17,7 @@ import {
   checkExecutionUsageLimits,
   USAGE_GATE_SETTLE_TIMEOUT_MS,
   USAGE_GATE_TTL_MS,
+  usageAdmissionTtlMs,
 } from '@/lib/billing/core/usage-gate-cache'
 import type { UsageUpgradePayer } from '@/lib/billing/usage-upgrade'
 import { coalesceLocally } from '@/lib/concurrency/singleflight'
@@ -141,7 +142,7 @@ export async function readMidRunUsageVerdict(
 }
 
 /**
- * Admitted direct-v1 verdicts, served for the execution gate's TTL like
+ * Admitted direct-v1 verdicts, served for {@link usageAdmissionTtlMs} like
  * {@link checkExecutionUsageLimits} serves attributed ones: the worker re-validates a run on
  * every resume leg, and each uncached read sums the payer's ledger for the period. Only a
  * `within` verdict is stored, so a refusal or an unreadable ledger is always read again.
@@ -227,7 +228,7 @@ export async function readMidRunAccountUsageVerdict(
       }
     }
     const within: MidRunUsageVerdict = { status: 'within' }
-    accountVerdictCache.set(key, within)
+    accountVerdictCache.set(key, within, { ttl: usageAdmissionTtlMs(usage) })
     return within
   } catch (error) {
     logger.warn('Mid-run account usage read failed; continuing the run', {

apps/sim/lib/billing/core/usage-gate-cache.test.ts

@@ -5,7 +5,7 @@ import {
   billingAttributionMockFns,
 } from '@sim/testing/mocks/billing-attribution.mock'
 import { sleep } from '@sim/utils/helpers'
-import { beforeEach, describe, expect, it, vi } from 'vitest'
+import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'
 
 vi.mock('@/lib/billing/core/attribution', () => billingAttributionMock)
 
@@ -14,13 +14,38 @@ import {
   checkExecutionUsageLimits,
   checkIngestionUsageLimits,
   checkSearchUsageLimits,
+  NEAR_LIMIT_USAGE_GATE_TTL_MS,
   resetUsageGateCache,
   USAGE_GATE_SETTLE_TIMEOUT_MS,
   USAGE_GATE_TTL_MS,
 } from '@/lib/billing/core/usage-gate-cache'
 
 const mockCheck = billingAttributionMockFns.mockCheckAttributedUsageLimits
 
+/** Fakes `setTimeout` for {@link advanceClockPast}; install it before the cache is first read. */
+function useFakeTimeouts(): void {
+  vi.useFakeTimers({ toFake: ['setTimeout', 'clearTimeout'] })
+}
+
+/**
+ * Fires the cache's pending clock debounce before restoring real timers; dropped unfired, it would
+ * pin the module-level cache to this test's spied clock for every later test.
+ */
+function restoreRealTimers(): void {
+  vi.runOnlyPendingTimers()
+  vi.useRealTimers()
+}
+
+/**
+ * `lru-cache` captures `performance` at load, so fake timers cannot move its clock: spy the clock
+ * forward, then fire the faked 1 ms timer behind which it debounces `performance.now()`.
+ */
+async function advanceClockPast(ms: number): Promise<void> {
+  const start = performance.now()
+  vi.spyOn(performance, 'now').mockReturnValue(start + ms + 1)
+  await vi.advanceTimersByTimeAsync(1)
+}
+
 const ATTRIBUTION: BillingAttributionSnapshot = {
   actorUserId: 'member-1',
   workspaceId: null,
@@ -74,17 +99,19 @@ describe('checkIngestionUsageLimits', () => {
   })
 
   it('keeps a refusal for the same bounded window, then reads fresh', async () => {
-    mockCheck.mockResolvedValueOnce({ isExceeded: true, scope: 'payer', message: 'over' })
-    expect((await checkIngestionUsageLimits(ATTRIBUTION)).isExceeded).toBe(true)
-    expect((await checkIngestionUsageLimits(ATTRIBUTION)).isExceeded).toBe(true)
-    expect(mockCheck).toHaveBeenCalledTimes(1)
+    useFakeTimeouts()
+    try {
+      mockCheck.mockResolvedValueOnce({ isExceeded: true, scope: 'payer', message: 'over' })
+      expect((await checkIngestionUsageLimits(ATTRIBUTION)).isExceeded).toBe(true)
+      expect((await checkIngestionUsageLimits(ATTRIBUTION)).isExceeded).toBe(true)
+      expect(mockCheck).toHaveBeenCalledTimes(1)
 
-    /** `lru-cache` reads `performance.now()` and debounces it behind a real 1 ms timer. */
-    const start = performance.now()
-    vi.spyOn(performance, 'now').mockReturnValue(start + USAGE_GATE_TTL_MS + 1)
-    await sleep(5)
-    expect((await checkIngestionUsageLimits(ATTRIBUTION)).isExceeded).toBe(false)
-    expect(mockCheck).toHaveBeenCalledTimes(2)
+      await advanceClockPast(USAGE_GATE_TTL_MS)
+      expect((await checkIngestionUsageLimits(ATTRIBUTION)).isExceeded).toBe(false)
+      expect(mockCheck).toHaveBeenCalledTimes(2)
+    } finally {
+      restoreRealTimers()
+    }
   })
 
   it('separates answers by actor, period, period source, payer and plan', async () => {
@@ -212,3 +239,72 @@ describe('checkExecutionUsageLimits', () => {
     }
   })
 })
+
+describe('admission lifetime by headroom', () => {
+  beforeEach(() => {
+    useFakeTimeouts()
+    resetUsageGateCache()
+    mockCheck.mockReset()
+  })
+
+  afterEach(() => {
+    restoreRealTimers()
+  })
+
+  it('re-reads a free-plan admission after the near-limit window, so spend past the limit is refused', async () => {
+    mockCheck
+      .mockResolvedValueOnce({ isExceeded: false, payerUsage: { currentUsage: 1, limit: 5 } })
+      .mockResolvedValueOnce({
+        isExceeded: true,
+        scope: 'payer',
+        message: 'over',
+        payerUsage: { currentUsage: 5.4, limit: 5 },
+      })
+    expect((await checkExecutionUsageLimits(ATTRIBUTION)).isExceeded).toBe(false)
+    expect((await checkExecutionUsageLimits(ATTRIBUTION)).isExceeded).toBe(false)
+    expect(mockCheck).toHaveBeenCalledTimes(1)
+
+    await advanceClockPast(NEAR_LIMIT_USAGE_GATE_TTL_MS)
+    expect((await checkExecutionUsageLimits(ATTRIBUTION)).isExceeded).toBe(true)
+    expect(mockCheck).toHaveBeenCalledTimes(2)
+  })
+
+  it('re-reads a large pool in the last fifth of its limit', async () => {
+    mockCheck.mockResolvedValue({
+      isExceeded: false,
+      payerUsage: { currentUsage: 900, limit: 1000 },
+    })
+    await checkExecutionUsageLimits(ATTRIBUTION)
+    await advanceClockPast(NEAR_LIMIT_USAGE_GATE_TTL_MS)
+    await checkExecutionUsageLimits(ATTRIBUTION)
+    expect(mockCheck).toHaveBeenCalledTimes(2)
+  })
+
+  it('re-reads when the acting member is near their cap though the pool has room', async () => {
+    mockCheck.mockResolvedValue({
+      isExceeded: false,
+      payerUsage: { currentUsage: 100, limit: 1000 },
+      memberUsage: { currentUsage: 48, limit: 50 },
+    })
+    await checkSearchUsageLimits(ATTRIBUTION)
+    await advanceClockPast(NEAR_LIMIT_USAGE_GATE_TTL_MS)
+    await checkSearchUsageLimits(ATTRIBUTION)
+    expect(mockCheck).toHaveBeenCalledTimes(2)
+  })
+
+  it('holds an admission with ample headroom for the full TTL', async () => {
+    mockCheck.mockResolvedValue({
+      isExceeded: false,
+      payerUsage: { currentUsage: 100, limit: 1000 },
+      memberUsage: { currentUsage: 10, limit: null },
+    })
+    await checkExecutionUsageLimits(ATTRIBUTION)
+    await advanceClockPast(NEAR_LIMIT_USAGE_GATE_TTL_MS)
+    await checkExecutionUsageLimits(ATTRIBUTION)
+    expect(mockCheck).toHaveBeenCalledTimes(1)
+
+    await advanceClockPast(USAGE_GATE_TTL_MS)
+    await checkExecutionUsageLimits(ATTRIBUTION)
+    expect(mockCheck).toHaveBeenCalledTimes(2)
+  })
+})

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

@@ -13,14 +13,50 @@ import { coalesceLocally } from '@/lib/concurrency/singleflight'
  * The gate sums the payer's usage ledger for the billing period, which grows
  * with the payer's activity, so a large organization scans its whole period on
  * every uncached call. Bulk ingestion re-checks per document and knowledge
- * search checks per query. Staleness is bounded by this TTL and fails in the
- * harmless direction: a payer who crosses their limit keeps going for at most
- * this long, which charges nobody wrongly. Five minutes: the sum is a few
- * hundred milliseconds for a busy payer, and a minute made every search after
- * a pause pay it.
+ * search checks per query. A payer who crosses their limit keeps going for at
+ * most this long, so only an admission with ample headroom is held this long
+ * (see {@link usageAdmissionTtlMs}). Five minutes: the sum is a few hundred
+ * milliseconds for a busy payer, and a minute made every search after a pause
+ * pay it.
  */
 export const USAGE_GATE_TTL_MS = 5 * 60 * 1000
 
+/**
+ * How long an admission near its limit stays usable. Every spend landing
+ * inside this window is spend past the limit nobody bills (a free plan) or
+ * past a cap the payer set, so it is short; it stays above zero so a burst of
+ * executions or model steps from one payer shares a read instead of summing
+ * the ledger on each. Staleness is this plus any lag in the usage read itself.
+ */
+export const NEAR_LIMIT_USAGE_GATE_TTL_MS = 10 * 1000
+
+/**
+ * An admission is near its limit when the remaining headroom is under the
+ * larger of this floor and {@link NEAR_LIMIT_HEADROOM_SHARE} of the limit. The
+ * floor puts every free-plan payer permanently near its limit; the share keeps
+ * a large pool from caching its last stretch.
+ */
+const NEAR_LIMIT_HEADROOM_FLOOR_USD = 25
+const NEAR_LIMIT_HEADROOM_SHARE = 0.2
+
+type UsageStanding = NonNullable<AttributedUsageLimitsResult['memberUsage']>
+
+function isNearLimit(usage: UsageStanding | undefined): boolean {
+  if (!usage || usage.limit === null) return false
+  const headroom = usage.limit - usage.currentUsage
+  return headroom < Math.max(NEAR_LIMIT_HEADROOM_FLOOR_USD, usage.limit * NEAR_LIMIT_HEADROOM_SHARE)
+}
+
+/**
+ * How long an admission may be served from cache: the full gate TTL while
+ * every standing (the payer's, the acting member's) has ample headroom,
+ * otherwise {@link NEAR_LIMIT_USAGE_GATE_TTL_MS}. An admission that carries no
+ * figures (billing off) has no limit to approach.
+ */
+export function usageAdmissionTtlMs(...standings: Array<UsageStanding | undefined>): number {
+  return standings.some(isNearLimit) ? NEAR_LIMIT_USAGE_GATE_TTL_MS : USAGE_GATE_TTL_MS
+}
+
 /**
  * How long a coalesced usage read may take before its callers give up on it. The read's cost is
  * its ledger aggregates, each of which the database ends at
@@ -78,7 +114,8 @@ function gateKey(attribution: BillingAttributionSnapshot): string {
  * `coalesceLocally` collapses concurrent misses onto one ledger read and bounds
  * a hung read at {@link USAGE_GATE_SETTLE_TIMEOUT_MS}. The write stays on the
  * value this caller received, so a producer that timed out and later resolved
- * cannot overwrite a fresher answer.
+ * cannot overwrite a fresher answer. An admission is stored for
+ * {@link usageAdmissionTtlMs}, a refusal for the full TTL.
  *
  * There is deliberately no invalidator: usage and limit changes land in other
  * processes (execution workers, Stripe webhooks), so the TTL is the real bound.
@@ -96,7 +133,13 @@ async function checkUsageLimitsThroughCache(
     () => checkAttributedUsageLimits(attribution),
     USAGE_GATE_SETTLE_TIMEOUT_MS
   )
-  if (cacheRefusals || !result.isExceeded) gateCache.set(key, result)
+  if (cacheRefusals || !result.isExceeded) {
+    gateCache.set(key, result, {
+      ttl: result.isExceeded
+        ? USAGE_GATE_TTL_MS
+        : usageAdmissionTtlMs(result.payerUsage, result.memberUsage),
+    })
+  }
   return result
 }
 

packages/testing/src/mocks/billing-usage-gate-cache.mock.ts

@@ -5,7 +5,7 @@ import { vi } from 'vitest'
  *
  * The three `check*UsageLimits` gates are bare `vi.fn()` (resolve `undefined`); set
  * `mockResolvedValue({ isExceeded: false })` for an admitted path. `resetUsageGateCache` is a
- * no-op.
+ * no-op, and `usageAdmissionTtlMs` returns the full gate TTL.
  *
  * @example
  * ```ts
@@ -21,6 +21,8 @@ export const billingUsageGateCacheMockFns = {
   mockResetUsageGateCache: vi.fn((): void => {}),
 }
 
+const USAGE_GATE_TTL_MS = 5 * 60 * 1000
+
 /**
  * Static mock module for `@/lib/billing/core/usage-gate-cache`. Constants carry the real
  * values (`USAGE_GATE_SETTLE_TIMEOUT_MS` = 2 x the 60 s ledger statement timeout + 15 s).
@@ -31,10 +33,11 @@ export const billingUsageGateCacheMockFns = {
  * ```
  */
 export const billingUsageGateCacheMock = {
-  USAGE_GATE_TTL_MS: 5 * 60 * 1000,
+  USAGE_GATE_TTL_MS,
   USAGE_GATE_SETTLE_TIMEOUT_MS: 2 * 60_000 + 15_000,
   checkIngestionUsageLimits: billingUsageGateCacheMockFns.mockCheckIngestionUsageLimits,
   checkSearchUsageLimits: billingUsageGateCacheMockFns.mockCheckSearchUsageLimits,
   checkExecutionUsageLimits: billingUsageGateCacheMockFns.mockCheckExecutionUsageLimits,
   resetUsageGateCache: billingUsageGateCacheMockFns.mockResetUsageGateCache,
+  usageAdmissionTtlMs: vi.fn((): number => USAGE_GATE_TTL_MS),
 }

scripts/test-patterns-baseline.json

@@ -151,7 +151,6 @@
   "real-wait\tapps/sim/hooks/queries/workspace.test.tsx\tsleep() #1",
   "real-wait\tapps/sim/lib/billing/core/reporting-usage-cache.test.ts\tsleep() #1",
   "real-wait\tapps/sim/lib/billing/core/reporting-usage-cache.test.ts\tsleep() #2",
-  "real-wait\tapps/sim/lib/billing/core/usage-gate-cache.test.ts\tsleep() #1",
   "real-wait\tapps/sim/lib/concurrency/leader-lock.test.ts\tsleep() #1",
   "real-wait\tapps/sim/lib/concurrency/singleflight.test.ts\tsleep() #1",
   "real-wait\tapps/sim/lib/core/async-jobs/backends/database.test.ts\tsleep() #1",