fix(billing): refuse expired reads and recover transaction timeouts

Author: waleedlatif1

apps/sim/lib/api/server/routes/billing-read.ts

@@ -4,14 +4,17 @@ import {
   internalErrorResponse,
   internalOrchestrationErrorPolicy,
 } from '@/lib/api/server/routes/internal-json-route'
-import { isTransientDatabaseReadError } from '@/lib/db/read-retry'
+import { DatabaseReadDeadlineError, isTransientDatabaseReadError } from '@/lib/db/read-retry'
 
 export const internalBillingReadErrorPolicy = extendInternalErrorPolicy(
   internalOrchestrationErrorPolicy,
   (error) => {
+    const cancellationReason = getPostgresCancellationReason(error)
     if (
+      !(error instanceof DatabaseReadDeadlineError) &&
       !isTransientDatabaseReadError(error) &&
-      getPostgresCancellationReason(error) !== 'statement_timeout'
+      cancellationReason !== 'statement_timeout' &&
+      cancellationReason !== 'transaction_timeout'
     )
       return null
     return internalErrorResponse(

apps/sim/lib/billing/application/organization-billing/get-organization-plan-seats.integration.ts

@@ -122,6 +122,8 @@ describe('organization plan and seats without a usage ledger', () => {
   it.each([
     { code: '40001', message: 'fixture recovery conflict', status: 503 },
     { code: '57014', message: 'canceling statement due to statement timeout', status: 503 },
+    { code: '25P04', message: 'terminating connection due to transaction timeout', status: 503 },
+    { code: '57014', message: 'canceling statement due to user request', status: 500 },
     { code: '42P01', message: 'fixture permanent database error', status: 500 },
   ])(
     'projects $code over HTTP without leaking database details',

apps/sim/lib/billing/core/ledger-read.integration.ts

@@ -7,10 +7,15 @@ import { generateId } from '@sim/utils/id'
 import { sql } from 'drizzle-orm'
 import { drizzle } from 'drizzle-orm/postgres-js'
 import postgres from 'postgres'
-import { afterAll, beforeAll, beforeEach, describe, expect, it } from 'vitest'
+import { afterAll, beforeAll, beforeEach, describe, expect, it, vi } from 'vitest'
 import { USAGE_LEDGER_STATEMENT_TIMEOUT_MS } from '@/lib/billing/constants'
 import { readLedgerBounded } from '@/lib/billing/core/ledger-read'
 
+vi.mock('@/lib/billing/constants', async (importOriginal) => ({
+  ...(await importOriginal<typeof import('@/lib/billing/constants')>()),
+  USAGE_LEDGER_STATEMENT_TIMEOUT_MS: 1000,
+}))
+
 const schemaName = `ledger_read_${generateId().replaceAll('-', '')}`
 const connection = postgres(readTestDatabaseUrl(), {
   max: 1,
@@ -79,6 +84,25 @@ describe('bounded ledger reads with PostgreSQL', () => {
     expect(await observeRead).toBeLessThanOrEqual(USAGE_LEDGER_STATEMENT_TIMEOUT_MS - 80)
   })
 
+  it('rejects a ledger read whose connection wait exhausted the budget', async () => {
+    let outcome: Promise<number> | undefined
+    await connection.begin(async () => {
+      outcome = readLedgerBounded(database, async (tx) => {
+        const rows = await tx.execute<{ total: string }>(
+          sql`select sum(cost) as total from ledger_probe`
+        )
+        return Number(rows[0].total)
+      })
+      void outcome.catch(() => undefined)
+      await sleep(USAGE_LEDGER_STATEMENT_TIMEOUT_MS + 100)
+    })
+    if (!outcome) throw new Error('Ledger read was not started')
+    await expect(outcome).rejects.toThrow('Database read deadline exceeded')
+    expect(await connection`select sum(cost)::text as total from ledger_probe`).toEqual([
+      { total: '1.00' },
+    ])
+  })
+
   it('cannot replay a write under the read-retry policy', async () => {
     const error = await readLedgerBounded(database, async (tx) => {
       await tx.execute(sql`insert into ledger_probe values (999)`)

apps/sim/lib/billing/core/ledger-read.ts

@@ -1,11 +1,12 @@
 import { sql } from 'drizzle-orm'
 import { USAGE_LEDGER_STATEMENT_TIMEOUT_MS } from '@/lib/billing/constants'
-import { withDatabaseReadRetry } from '@/lib/db/read-retry'
+import { DatabaseReadDeadlineError, withDatabaseReadRetry } from '@/lib/db/read-retry'
 import type { DbClient, DbTransaction } from '@/lib/db/types'
 
 /**
- * Runs one aggregate over a payer's usage ledger in a transaction of its own, bounded by
- * {@link USAGE_LEDGER_STATEMENT_TIMEOUT_MS}. `SET LOCAL` scopes the bound to that transaction,
+ * Runs one aggregate over a payer's usage ledger in a transaction of its own, sharing
+ * {@link USAGE_LEDGER_STATEMENT_TIMEOUT_MS} across retries and statement execution. Pool wait
+ * is deducted on acquisition; an exhausted budget refuses the read. `SET LOCAL` scopes the bound,
  * so it ends with the read and never reaches the pool. Every sum over a payer's billing period
  * reads through here, whether it admits a run, closes a cycle or previews a bill: a payer whose
  * period has grown past what one statement can sum within the bound fails at the database
@@ -22,12 +23,10 @@ export function readLedgerBounded<T>(
       const attemptStartedAt = Date.now()
       return executor.transaction(
         async (tx) => {
-          const timeoutMs = Math.max(
-            1,
-            Math.floor(
-              (remainingMs ?? USAGE_LEDGER_STATEMENT_TIMEOUT_MS) - (Date.now() - attemptStartedAt)
-            )
+          const timeoutMs = Math.floor(
+            (remainingMs ?? USAGE_LEDGER_STATEMENT_TIMEOUT_MS) - (Date.now() - attemptStartedAt)
           )
+          if (timeoutMs <= 0) throw new DatabaseReadDeadlineError()
           await tx.execute(sql.raw(`SET LOCAL statement_timeout = '${timeoutMs}ms'`))
           return read(tx)
         },

apps/sim/lib/db/read-retry.ts

@@ -27,6 +27,14 @@ const TRANSIENT_READ_CODES = new Set([
   'ETIMEDOUT',
 ])
 
+/** A read exhausted its budget before its query could begin. */
+export class DatabaseReadDeadlineError extends Error {
+  constructor() {
+    super('Database read deadline exceeded')
+    this.name = 'DatabaseReadDeadlineError'
+  }
+}
+
 /** Only driver codes identify retryable reads; query cancellation and application errors propagate. */
 export function isTransientDatabaseReadError(error: unknown): boolean {
   const code = getPostgresErrorCode(error)