fix(billing): cache admitted direct-v1 continuation verdicts and bound the callback's standing read

- A direct-v1 continuation re-read the payer's full period ledger on every resume leg. Its
  admitted verdict is now served for the execution gate's TTL, with concurrent misses coalesced,
  like the attributed path; a refusal or an unreadable ledger is always read again, and the read
  is skipped when billing is off.
- The cost callback waits at most 1 s on the payer's standing, well inside the worker's 5 s
  callback timeout, and answers not exceeded past it; the abandoned read still caches its
  admission.
- The straddling-period test now reaches the re-judge branch.

Author: waleedlatif1

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

@@ -49,7 +49,7 @@ vi.mock('@/lib/billing/threshold-billing', () => ({
 }))
 
 import { billingUpdateCostResponseSchema } from '@/lib/api/contracts/subscription'
-import { resetMidRunPeriodCache } from '@/lib/billing/core/mid-run-usage'
+import { resetMidRunUsageCaches } from '@/lib/billing/core/mid-run-usage'
 import { resetUsageGateCache } from '@/lib/billing/core/usage-gate-cache'
 import {
   BillingCallbackBody,
@@ -897,7 +897,7 @@ describe('POST /api/billing/update-cost — mid-run usage gate', () => {
 
   beforeEach(() => {
     resetUsageGateCache()
-    resetMidRunPeriodCache()
+    resetMidRunUsageCaches()
     setEnvFlags({ isBillingEnabled: true, isHosted: true })
     mockCheckInternalApiKey.mockReturnValue({ success: true })
     mockRecordCumulativeUsage.mockResolvedValue({ billed: true, delta: 0.5, total: 0.5 })
@@ -1136,8 +1136,9 @@ describe('POST /api/billing/update-cost — mid-run usage gate', () => {
         end: new Date(Date.now() + 40).toISOString(),
       },
     }
-    mockRequireBillingAttributionHeader.mockReturnValue(straddling)
-    mockRefreshAttributionPeriod.mockResolvedValue(CURRENT_ATTRIBUTION)
+    mockRefreshAttributionPeriod
+      .mockResolvedValueOnce(straddling)
+      .mockResolvedValue(CURRENT_ATTRIBUTION)
     mockCheckAttributedUsageLimits.mockImplementation(
       async (attribution: typeof CURRENT_ATTRIBUTION) => {
         if (attribution.billingPeriod.end !== straddling.billingPeriod.end) {
@@ -1151,6 +1152,21 @@ describe('POST /api/billing/update-cost — mid-run usage gate', () => {
     const body = await (await POST(attributedCallback())).json()
 
     expect(body.usageExceeded).toBe(false)
+    expect(mockRefreshAttributionPeriod).toHaveBeenCalledTimes(2)
+  })
+
+  it('answers not exceeded when the standing read outlasts the callback budget', async () => {
+    mockCheckAttributedUsageLimits.mockImplementation(async () => {
+      await sleep(1500)
+      return { isExceeded: true, scope: 'payer' }
+    })
+    const startedAt = Date.now()
+
+    const res = await POST(attributedCallback())
+
+    expect(res.status).toBe(200)
+    await expect(res.json()).resolves.toMatchObject({ success: true, usageExceeded: false })
+    expect(Date.now() - startedAt).toBeLessThan(1400)
   })
 
   it('reloads a cached current period once it has ended', async () => {

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

@@ -22,7 +22,7 @@ import {
   resolveLegacyV0BillingAttribution,
   toBillingContext,
 } from '@/lib/billing/core/billing-attribution'
-import { readMidRunUsageVerdict } from '@/lib/billing/core/mid-run-usage'
+import { type MidRunUsageVerdict, readMidRunUsageVerdict } from '@/lib/billing/core/mid-run-usage'
 import {
   type CumulativeUsageContextField,
   CumulativeUsageContextMismatchError,
@@ -35,6 +35,7 @@ import {
 } from '@/lib/billing/threshold-billing'
 import { resolveUsageUpgradePayload } from '@/lib/billing/usage-upgrade'
 import { isBillingEnabled, isHosted } from '@/lib/core/config/env-flags'
+import { withinDeadline } from '@/lib/core/utils/deadline'
 import { generateRequestId } from '@/lib/core/utils/request'
 import { withRouteHandler } from '@/lib/core/utils/with-route-handler'
 import { BILLING_CALLBACK_OUTCOME } from '@/lib/mothership/generated/billing-protocol-v1'
@@ -45,6 +46,14 @@ import { checkInternalApiKey } from '@/lib/mothership/request/http'
 import { withIncomingGoSpan } from '@/lib/mothership/request/otel'
 
 const logger = createLogger('BillingUpdateCostAPI')
+/**
+ * How long a cost callback waits on the payer's standing. The worker gives up on the whole
+ * callback after 5 s, and a cold gate read can wait on the ledger far longer; past this the
+ * callback answers not-exceeded. The abandoned read keeps running and caches its admission, and
+ * the next step or re-check reads a refusal again.
+ */
+const USAGE_STANDING_TIMEOUT_MS = 1000
+
 const RETRYABLE_SETTLEMENT_RESPONSE = {
   code: 'BILLING_SETTLEMENT_RETRYABLE',
   error: 'Billing settlement temporarily unavailable',
@@ -70,14 +79,24 @@ function invalidBillingProtocolResponse(requestId: string, span: Span): NextResp
  * Served from the execution usage gate: 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
  * 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.
+ * refusal to the next step or re-check rather than ending a paying run on a database blip,
+ * and so does a read that outlasts {@link USAGE_STANDING_TIMEOUT_MS}.
  */
 async function readUsageStanding(
   userId: string,
   billingAttribution: BillingAttributionSnapshot | undefined
 ): Promise<BillingUsageVerdict> {
   if (!isHosted || !billingAttribution) return { usageExceeded: false }
-  const verdict = await readMidRunUsageVerdict(billingAttribution)
+  let verdict: MidRunUsageVerdict
+  try {
+    verdict = await withinDeadline(
+      () => readMidRunUsageVerdict(billingAttribution),
+      Date.now() + USAGE_STANDING_TIMEOUT_MS
+    )
+  } catch {
+    logger.warn('Usage standing read outlasted the callback budget; answering not exceeded')
+    return { usageExceeded: false }
+  }
   // Only a spent limit pauses the run. A blocked account is refused at the run's next
   // continuation or re-check, with blocked-account copy rather than the upgrade card.
   if (verdict.status !== 'exceeded') return { usageExceeded: false }

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

@@ -141,7 +141,7 @@ vi.mock('@/lib/workspaces/permissions/utils', () => permissionsMock)
 vi.mock('@/lib/workspaces/utils', () => workspacesUtilsMock)
 
 import { validateCopilotApiKeyBodySchema } from '@/lib/api/contracts/copilot'
-import { resetMidRunPeriodCache } from '@/lib/billing/core/mid-run-usage'
+import { resetMidRunUsageCaches } from '@/lib/billing/core/mid-run-usage'
 import { resetUsageGateCache } from '@/lib/billing/core/usage-gate-cache'
 import { POST } from '@/app/api/copilot/api-keys/validate/route'
 
@@ -524,7 +524,7 @@ describe('validation lifecycle purposes', () => {
       periodEnd: new Date(ATTRIBUTION.billingPeriod.end),
     })
     resetUsageGateCache()
-    resetMidRunPeriodCache()
+    resetMidRunUsageCaches()
   })
 
   it('defaults older callers to full admission and rejects unknown purposes', () => {
@@ -796,6 +796,28 @@ describe('validation lifecycle purposes', () => {
     expect(response.status).toBe(402)
   })
 
+  it('answers repeated direct-v1 continuations from the cached admission and re-reads a refusal', async () => {
+    for (let call = 0; call < 2; call++) queueTableRows(schemaMock.user, [{ id: 'user-1' }])
+    for (let leg = 0; leg < 3; leg++) {
+      expect((await POST(request(body, directHeaders))).status).toBe(200)
+    }
+    expect(mockCheckUsageStatus).toHaveBeenCalledTimes(1)
+
+    resetMidRunUsageCaches()
+    mockCheckUsageStatus.mockResolvedValue({ isExceeded: true, currentUsage: 12, limit: 10 })
+    expect((await POST(request(body, directHeaders))).status).toBe(402)
+    expect((await POST(request(body, directHeaders))).status).toBe(402)
+    expect(mockCheckUsageStatus).toHaveBeenCalledTimes(3)
+  })
+
+  it('never reads the ledger for a direct-v1 continuation when billing is off', async () => {
+    setEnvFlags({ isHosted: false, isBillingEnabled: false })
+    mockCheckUsageStatus.mockResolvedValue({ isExceeded: true, currentUsage: 12, limit: 10 })
+
+    expect((await POST(request(body, directHeaders))).status).toBe(200)
+    expect(mockCheckUsageStatus).not.toHaveBeenCalled()
+  })
+
   it('judges a direct-v1 organization payer without a subscription as that organization', async () => {
     mockGetOrganizationSubscription.mockResolvedValue(null)
     mockCheckUsageStatus.mockImplementation(

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

@@ -12,7 +12,13 @@ import {
 import { defaultBillingPeriod } from '@/lib/billing/core/billing-period'
 import { getHighestPriorityPersonalSubscription } from '@/lib/billing/core/plan'
 import { resolveSubscriptionUsagePeriod } from '@/lib/billing/core/reporting-period'
-import { checkExecutionUsageLimits } from '@/lib/billing/core/usage-gate-cache'
+import {
+  checkExecutionUsageLimits,
+  USAGE_GATE_SETTLE_TIMEOUT_MS,
+  USAGE_GATE_TTL_MS,
+} from '@/lib/billing/core/usage-gate-cache'
+import { coalesceLocally } from '@/lib/concurrency/singleflight'
+import { isBillingEnabled, isHosted } from '@/lib/core/config/env-flags'
 
 const logger = createLogger('MidRunUsage')
 
@@ -121,6 +127,17 @@ export async function readMidRunUsageVerdict(
   return { status: 'unknown' }
 }
 
+/**
+ * Admitted direct-v1 verdicts, served for the execution gate's TTL 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.
+ */
+const accountVerdictCache = new LRUCache<string, MidRunUsageVerdict>({
+  max: 10_000,
+  ttl: USAGE_GATE_TTL_MS,
+})
+
 /**
  * The same verdict for a direct-v1 run billed to an account decision rather than an attributed
  * payer. The payer is the one saved in the decision at admission, never re-selected from the
@@ -129,6 +146,7 @@ export async function readMidRunUsageVerdict(
 export async function readMidRunAccountUsageVerdict(
   decision: AccountBillingDecision
 ): Promise<MidRunUsageVerdict> {
+  if (!isHosted || !isBillingEnabled) return { status: 'within' }
   try {
     const payer = decision.billingEntity
     const subscription =
@@ -139,6 +157,20 @@ export async function readMidRunAccountUsageVerdict(
       ...defaultBillingPeriod(),
       source: 'default' as const,
     }
+    const key = [
+      payer.type,
+      payer.id,
+      billingPeriod.start.toISOString(),
+      billingPeriod.end.toISOString(),
+      billingPeriod.source,
+      decision.userId,
+      subscription?.id ?? '',
+      subscription?.plan ?? '',
+      subscription?.status ?? '',
+      subscription?.seats ?? '',
+    ].join(':')
+    const cached = accountVerdictCache.get(key)
+    if (cached) return cached
     // An organization payer without a subscription stays organization-scoped on the free plan,
     // as `toUsageLimitSubscription` does for attributed runs, never the actor's personal ledger.
     const usageSubscription =
@@ -153,12 +185,20 @@ export async function readMidRunAccountUsageVerdict(
             periodEnd: billingPeriod.end,
           }
         : null)
-    const usage = await checkUsageStatus(decision.userId, usageSubscription, {
-      billingEntity: payer,
-      billingPeriod,
-    })
+    const usage = await coalesceLocally(
+      `mid-run-account-usage:${key}`,
+      () =>
+        checkUsageStatus(decision.userId, usageSubscription, {
+          billingEntity: payer,
+          billingPeriod,
+        }),
+      USAGE_GATE_SETTLE_TIMEOUT_MS
+    )
     if (usage.unavailable) return { status: 'unknown' }
-    return usage.isExceeded ? { status: 'exceeded', scope: 'payer' } : { status: 'within' }
+    if (usage.isExceeded) return { status: 'exceeded', scope: 'payer' }
+    const within: MidRunUsageVerdict = { status: 'within' }
+    accountVerdictCache.set(key, within)
+    return within
   } catch (error) {
     logger.warn('Mid-run account usage read failed; continuing the run', {
       error: getErrorMessage(error),
@@ -167,7 +207,8 @@ export async function readMidRunAccountUsageVerdict(
   }
 }
 
-/** Drops every cached current period. Test seam; never called in production code. */
-export function resetMidRunPeriodCache(): void {
+/** Drops every cached current period and account verdict. Test seam; never called in production code. */
+export function resetMidRunUsageCaches(): void {
   currentPeriodCache.clear()
+  accountVerdictCache.clear()
 }

apps/sim/lib/mothership/request/lifecycle/admission.test.ts

@@ -6,7 +6,7 @@ import {
 } from '@sim/testing/mocks/billing-usage-gate-cache.mock'
 import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'
 import { createAttributedBillingRequestEnvelope } from '@/lib/billing/core/billing-attribution'
-import { resetMidRunPeriodCache } from '@/lib/billing/core/mid-run-usage'
+import { resetMidRunUsageCaches } from '@/lib/billing/core/mid-run-usage'
 import { OrchestrationError } from '@/lib/core/orchestration/types'
 import { BillingLimitError } from '@/lib/mothership/request/go/stream'
 import { authorizeLifecycleContinuation, restoreBillingAdmission } from './admission'
@@ -49,7 +49,7 @@ beforeEach(() => {
     periodStart: new Date(attribution.billingPeriod.start),
     periodEnd: new Date(attribution.billingPeriod.end),
   })
-  resetMidRunPeriodCache()
+  resetMidRunUsageCaches()
 })
 afterEach(resetEnvFlagsMock)
 
@@ -157,7 +157,7 @@ describe('continuation admission', () => {
       authorizeLifecycleContinuation({ ...context, billingAttribution: ended })
     ).rejects.toBeInstanceOf(BillingLimitError)
 
-    resetMidRunPeriodCache()
+    resetMidRunUsageCaches()
     mockGetOrganizationSubscription.mockRejectedValue(new Error('subscription read failed'))
     await expect(
       authorizeLifecycleContinuation({ ...context, billingAttribution: ended })