fix(billing): answer an in-flight retry at once, and tell Chat what provenance withheld

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

@@ -1,5 +1,6 @@
 import { readFileSync } from 'node:fs'
 import { createMockRequest, resetEnvFlagsMock, setEnvFlags } from '@sim/testing'
+import { createRouteContext } from '@sim/testing/helpers/http'
 import {
   billingAttributionMock,
   billingAttributionMockFns,
@@ -9,6 +10,7 @@ import { billingPlanMock, billingPlanMockFns } from '@sim/testing/mocks/billing-
 import {
   billingUsageLogMock,
   billingUsageLogMockFns,
+  MockCumulativeUsageInFlightError,
 } from '@sim/testing/mocks/billing-usage-log.mock'
 import {
   billingUsageMonitorMock,
@@ -18,6 +20,7 @@ import { copilotHttpMock, copilotHttpMockFns } from '@sim/testing/mocks/copilot-
 import { mothershipOtelMock } from '@sim/testing/mocks/mothership-otel.mock'
 import { sleep } from '@sim/utils/helpers'
 import { afterAll, afterEach, beforeEach, describe, expect, it, vi } from 'vitest'
+import { BillingInFlightCode } from '@/lib/mothership/generated/billing'
 
 const {
   mockCheckAndBillOverageThreshold,
@@ -748,6 +751,42 @@ describe('POST /api/billing/update-cost — workspaceId attribution', () => {
     })
   })
 
+  it('answers a retryable in-flight 503 while another attempt of the request is recorded', async () => {
+    const billingRequestId = '0190c03f-9f7d-4b79-8b58-e7f779fd29e1'
+    mockRecordCumulativeUsage.mockRejectedValueOnce(
+      new MockCumulativeUsageInFlightError(`update-cost:${billingRequestId}`)
+    )
+
+    const res = await POST(
+      createMockRequest(
+        'POST',
+        {
+          userId: 'user-1',
+          cost: 0.5,
+          model: 'claude-opus-4.8',
+          source: 'workspace-chat',
+          idempotencyKey: billingRequestId,
+        },
+        {
+          'x-api-key': 'internal',
+          'x-sim-billing-protocol': 'direct-v1',
+          'x-sim-billing-request-id': billingRequestId,
+          'x-sim-billing-account-decision': 'serialized-account-decision',
+        }
+      ),
+      createRouteContext({})
+    )
+
+    expect(res.status).toBe(503)
+    expect(res.headers.get('retry-after')).toBe('5')
+    await expect(res.json()).resolves.toMatchObject({
+      success: false,
+      code: BillingInFlightCode,
+      retryable: true,
+    })
+    expect(mockCheckAndBillPayerOverageThreshold).not.toHaveBeenCalled()
+  })
+
   it('retries modern settlement on a duplicate cumulative callback before returning 409', async () => {
     const billingRequestId = '0190c03f-9f7d-4b79-8b58-e7f779fd29e1'
     mockRecordCumulativeUsage

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

@@ -30,6 +30,7 @@ import {
 import {
   type CumulativeUsageContextField,
   CumulativeUsageContextMismatchError,
+  CumulativeUsageInFlightError,
   CumulativeUsagePeriodClosedError,
   recordCumulativeUsage,
 } from '@/lib/billing/core/usage-log'
@@ -43,6 +44,7 @@ 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 { BillingInFlightCode } from '@/lib/mothership/generated/billing'
 import { BILLING_CALLBACK_OUTCOME } from '@/lib/mothership/generated/billing-protocol-v1'
 import { BillingRouteOutcome } from '@/lib/mothership/generated/trace-attribute-values-v1'
 import { TraceAttr } from '@/lib/mothership/generated/trace-attributes-v1'
@@ -53,7 +55,7 @@ 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 after 15 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.
  */
@@ -510,6 +512,32 @@ async function updateCostInner(req: NextRequest, span: Span): Promise<NextRespon
       )
     }
 
+    /**
+     * Another attempt of this request is still being recorded, so this one's outcome is unknown:
+     * the caller resends later, and the monotonic top-up applies the higher cumulative once.
+     */
+    if (error instanceof CumulativeUsageInFlightError) {
+      logger.info(`[${requestId}] Cost update already in flight for this request`, {
+        eventKey: error.eventKey,
+        duration,
+      })
+      // The closed outcome set has no in-flight value, so the span matches the retryable settlement
+      // 503; the body's BILLING_REQUEST_IN_FLIGHT code and the INFO line above tell them apart.
+      span.setAttribute(TraceAttr.BillingOutcome, BillingRouteOutcome.InternalError)
+      span.setAttribute(TraceAttr.HttpStatusCode, 503)
+      span.setAttribute(TraceAttr.BillingDurationMs, duration)
+      return NextResponse.json(
+        {
+          success: false,
+          code: BillingInFlightCode,
+          error: 'Billing request already in flight',
+          retryable: true,
+          requestId,
+        },
+        { status: 503, headers: { 'Retry-After': '5' } }
+      )
+    }
+
     const pgCode = getPostgresErrorCode(error)
     const pgConstraint = getPostgresConstraintName(error)
     const reconciliationOutcome =

apps/sim/lib/billing/core/usage-log.integration.ts

@@ -28,6 +28,7 @@ vi.mock('@/lib/billing/subscriptions/utils', async (importOriginal) => ({
 
 import {
   CumulativeUsageContextMismatchError,
+  CumulativeUsageInFlightError,
   CumulativeUsagePeriodClosedError,
   getBillingPeriodUsageCost,
   getBillingPeriodUsageCostByUser,
@@ -92,6 +93,18 @@ function usage(cost: number, eventKey = 'update-cost:shared-request'): RecordCum
   }
 }
 
+/** Resends a callback refused as in flight, the way the worker's replay does, until it is recorded. */
+async function untilRecorded(params: RecordCumulativeUsageParams) {
+  for (let attempt = 0; ; attempt++) {
+    try {
+      return await recordCumulativeUsage(params)
+    } catch (error) {
+      if (!(error instanceof CumulativeUsageInFlightError) || attempt >= 200) throw error
+      await sleep(50)
+    }
+  }
+}
+
 async function ledgerRows() {
   return connection<{ event_key: string; cost: string }[]>`
     select event_key, cost from usage_log order by event_key
@@ -199,13 +212,15 @@ describe('Cumulative billing with PostgreSQL', () => {
     }
   )
 
-  it('reproduces a retry timing out behind an idle holder when only lock waits are bounded', async () => {
+  it('refuses a retry behind a holder of the same request at once, without waiting on its lock', async () => {
     const pause = pauseNextTransaction(true)
     const holder = recordCumulativeUsage(usage(0.4))
     try {
       await pause.reached.promise
+      const startedAt = performance.now()
       const failure = await recordCumulativeUsage(usage(0.8)).catch((error: unknown) => error)
-      expect(getPostgresErrorCode(failure)).toBe('55P03')
+      expect(failure).toBeInstanceOf(CumulativeUsageInFlightError)
+      expect(performance.now() - startedAt).toBeLessThan(1_000)
       expect(await ledgerRows()).toEqual([])
     } finally {
       pause.release.resolve()
@@ -223,8 +238,8 @@ describe('Cumulative billing with PostgreSQL', () => {
       try {
         await pause.reached.promise
         const failure = await recordCumulativeUsage(usage(0.8)).catch((error: unknown) => error)
-        expect(getPostgresErrorCode(failure)).toBe('55P03')
-        const recovered = await recordCumulativeUsage(usage(0.8))
+        expect(failure).toBeInstanceOf(CumulativeUsageInFlightError)
+        const recovered = await untilRecorded(usage(0.8))
         expect(recovered.billed).toBe(true)
         expect(recovered.delta).toBeCloseTo(0.8 - initial, 9)
         expect(recovered.total).toBe(0.8)
@@ -249,7 +264,7 @@ describe('Cumulative billing with PostgreSQL', () => {
 
   it('converges concurrent out-of-order callbacks and independent events to their exact totals', async () => {
     const costs = [0.4, 0.1, 0.8, 0.3, 0.8, 0.6]
-    const results = await Promise.all(costs.map((cost) => recordCumulativeUsage(usage(cost))))
+    const results = await Promise.all(costs.map((cost) => untilRecorded(usage(cost))))
     expect(results.reduce((total, result) => total + result.delta, 0)).toBeCloseTo(0.8, 9)
     await Promise.all(
       Array.from({ length: 32 }, (_, index) =>

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

@@ -5,6 +5,7 @@ import {
   billingSubscriptionUtilsMock,
   billingSubscriptionUtilsMockFns,
 } from '@sim/testing/mocks/billing-subscription-utils.mock'
+import { getMockLogger } from '@sim/testing/mocks/logger.mock'
 import { afterAll, afterEach, beforeAll, beforeEach, describe, expect, it, vi } from 'vitest'
 
 const {
@@ -31,6 +32,7 @@ import { USAGE_LEDGER_STATEMENT_TIMEOUT_MS } from '@/lib/billing/constants'
 import {
   CUMULATIVE_COST_EPSILON,
   CumulativeUsageContextMismatchError,
+  CumulativeUsageInFlightError,
   getBillingPeriodUsageCost,
   getBillingPeriodUsageCostByUser,
   getBillingPeriodUsageCostWithSourceSubset,
@@ -245,7 +247,8 @@ describe('recordCumulativeUsage', () => {
     mockIsOrgScopedSubscription.mockReturnValue(true)
   })
 
-  const setupTx = (existingRow: Partial<typeof defaultExistingRow> | null) => {
+  /** `locked`: another attempt of the request holds its ledger lock. */
+  const setupTx = (existingRow: Partial<typeof defaultExistingRow> | null, locked = false) => {
     const resolvedExistingRow = existingRow ? { ...defaultExistingRow, ...existingRow } : null
     const limit = vi.fn().mockResolvedValue(resolvedExistingRow ? [resolvedExistingRow] : [])
     const where = vi.fn().mockReturnValue({ limit })
@@ -255,7 +258,8 @@ describe('recordCumulativeUsage', () => {
     const updateSet = vi.fn().mockReturnValue({ where: updateWhere })
     mockUpdate.mockReturnValue({ set: updateSet })
     const tx = {
-      execute: vi.fn().mockResolvedValue(undefined),
+      // The configure statement's result is unused; the try-lock reads `acquired`.
+      execute: vi.fn().mockResolvedValue([{ acquired: !locked }]),
       select,
       update: mockUpdate,
       insert: mockInsert, // recordUsage(tx) reuses the shared insert chain
@@ -290,6 +294,27 @@ describe('recordCumulativeUsage', () => {
     })
   })
 
+  it('refuses at once, writing nothing, while another attempt of the request holds its lock', async () => {
+    const { select } = setupTx(null, true)
+    await expect(
+      recordCumulativeUsage({
+        userId: 'user-1',
+        source: 'workspace-chat',
+        model: 'claude-opus-4.8',
+        cost: 0.4662453,
+        eventKey: 'update-cost:msg-1-billing',
+      })
+    ).rejects.toBeInstanceOf(CumulativeUsageInFlightError)
+    expect(select).not.toHaveBeenCalled()
+    expect(mockInsert).not.toHaveBeenCalled()
+    expect(mockUpdate).not.toHaveBeenCalled()
+    // The route logs the refusal once; it is not a transaction that failed to complete.
+    expect(getMockLogger('UsageLog').warn).not.toHaveBeenCalledWith(
+      'Cumulative usage transaction did not complete promptly',
+      expect.anything()
+    )
+  })
+
   it('tops up to the higher cumulative and bills only the delta', async () => {
     const { updateSet } = setupTx({ id: 'row-1', cost: '0.3474447' })
     const result = await recordCumulativeUsage({

apps/sim/lib/billing/core/usage-log.ts

@@ -30,7 +30,7 @@ import { isOrgScopedSubscription } from '@/lib/billing/subscriptions/utils'
 import type { InternalUsageLogSource } from '@/lib/billing/usage-sources'
 import { asOrchestrationError, OrchestrationError } from '@/lib/core/orchestration/types'
 import { HttpError } from '@/lib/core/utils/http-error'
-import { acquireAdvisoryXactLock } from '@/lib/db/advisory-locks'
+import { tryAcquireAdvisoryXactLock } from '@/lib/db/advisory-locks'
 import type { DbClient, DbOrTx } from '@/lib/db/types'
 
 const logger = createLogger('UsageLog')
@@ -700,6 +700,18 @@ export class CumulativeUsagePeriodClosedError extends Error {
   }
 }
 
+/**
+ * Another attempt of the same request holds its ledger lock, so this one cannot know whether that
+ * attempt's charge will land. Waiting would only hold a connection for an answer the caller gets by
+ * retrying later: the monotonic top-up applies whichever cumulative arrives once.
+ */
+export class CumulativeUsageInFlightError extends Error {
+  constructor(readonly eventKey: string) {
+    super(`Cumulative usage event "${eventKey}" is already being recorded`)
+    this.name = 'CumulativeUsageInFlightError'
+  }
+}
+
 interface CumulativeUsageLedgerBinding {
   userId: string
   workspaceId: string | null
@@ -755,11 +767,13 @@ function assertCumulativeUsageLedgerBinding(
 }
 
 /**
- * PostgreSQL 17+ bounds the entire transaction below the callback's five-second
- * deadline. Older supported servers instead bound each idle interval between
- * statements, alongside the per-statement budget. Both policies release an idle
- * lock holder without waiting for its application process to resume; only the
- * newer policy also limits total elapsed transaction time.
+ * PostgreSQL 17+ bounds the transaction's statements and the idle time between
+ * them; older supported servers bound each idle interval instead, alongside the
+ * per-statement budget. Both release an idle lock holder without waiting for its
+ * application process to resume. Neither bounds COMMIT: a commit waiting for a
+ * synchronous standby keeps its locks, including the request's ledger lock, until
+ * the standby confirms, however long that takes. `lock_timeout` bounds the row
+ * locks taken under the ledger lock; the ledger lock itself is never waited for.
  */
 const CUMULATIVE_FLUSH_TRANSACTION_TIMEOUT_MS = 4_000
 const CUMULATIVE_FLUSH_STATEMENT_TIMEOUT_MS = 3_500
@@ -778,10 +792,12 @@ type CumulativeUsageStage = 'pool' | 'configure' | 'lock' | 'read' | 'write' | '
  * (no two flushes can both believe they are first and clobber each other).
  * An existing row must match the incoming actor, workspace, payer, and billing
  * period before either a duplicate no-op or a top-up is accepted.
- * The billing context is resolved BEFORE the transaction and the lock wait is
- * bounded by `lock_timeout`. A server-enforced transaction deadline, or idle
- * transaction deadline on older PostgreSQL, releases a stalled holder. The
- * critical section uses one SELECT plus one INSERT/UPDATE on a single connection.
+ * The billing context is resolved BEFORE the transaction. The lock is taken
+ * without waiting: a held lock means another attempt of the same request is
+ * still being recorded (possibly committing behind a synchronous standby), and
+ * this attempt throws {@link CumulativeUsageInFlightError} instead of holding a
+ * connection while it waits. The critical section uses one SELECT plus one
+ * INSERT/UPDATE on a single connection.
  *
  * Because every leg flushes its cumulative and this converges to the max,
  * there is no under-billing if the request recovers after a partial flush, no
@@ -816,6 +832,7 @@ export async function recordCumulativeUsage(
   const stageDurationsMs: Partial<Record<CumulativeUsageStage, number>> = {}
   let succeeded = false
   let pgCode: string | undefined
+  let inFlight = false
   let toppedUpRowCreatedAt: Date | undefined
   const enterStage = (nextStage: CumulativeUsageStage) => {
     const now = Date.now()
@@ -841,7 +858,9 @@ export async function recordCumulativeUsage(
           set_config('lock_timeout', ${`${CUMULATIVE_FLUSH_LOCK_TIMEOUT_MS}ms`}, true)
       `)
       enterStage('lock')
-      await acquireAdvisoryXactLock(tx, 'usage_log_event', eventKey)
+      if (!(await tryAcquireAdvisoryXactLock(tx, 'usage_log_event', eventKey))) {
+        throw new CumulativeUsageInFlightError(eventKey)
+      }
 
       enterStage('read')
       const rows = await tx
@@ -997,12 +1016,14 @@ export async function recordCumulativeUsage(
     return result
   } catch (error) {
     pgCode = getPostgresErrorCode(error)
+    // Another attempt holding the lock is an answer, not a slow transaction; the route logs it.
+    if (error instanceof CumulativeUsageInFlightError) inFlight = true
     throw error
   } finally {
     const now = Date.now()
     stageDurationsMs[stage] = now - stageStartedAt
     const durationMs = now - startedAt
-    if (!succeeded || durationMs >= CUMULATIVE_FLUSH_SLOW_MS) {
+    if (!inFlight && (!succeeded || durationMs >= CUMULATIVE_FLUSH_SLOW_MS)) {
       logger.warn('Cumulative usage transaction did not complete promptly', {
         eventKey,
         succeeded,