fix(db): isolate billing and auth connection pools, release the redeploy-check transaction early, cache integration tool schemas (#9347)

* fix(db): give billing and auth their own connection pools

postgres.js pipelines new queries onto busy connections when none is idle,
so slow reads saturating the process pool made short usage-ledger
transactions wait seconds for BEGIN until Postgres killed them.

- Add `billing` (4) and `auth` (3) dbFor roles, with DATABASE_URL_BILLING
  and DATABASE_URL_AUTH overrides falling back to the process URL.
- Billing pool: recordCumulativeUsage, recordUsage without a caller tx,
  the execution logger's usage reconciliation, both threshold settlement
  transactions, and recordJobRunSpend. Each moved transaction only queries
  through its tx handle; readers stay on the process pool.
- Auth pool: getAuthPool() backs the Better Auth adapter, getAuthDatabase()
  outside a transaction, the OAuth adapter guard defaults, and the API-key
  hash lookup.
- isPrimaryPoolClient() in @sim/db replaces the `executor === db` checks in
  the security, session, and SSO policy caches, so reads through the auth
  pool keep hitting them instead of costing two queries per cached session.
- Integration test saturates the process pool and requires the billing
  write, an adapter session read, and an API-key lookup to finish anyway.

* fix(workflows): compare deployment states after the snapshot transaction commits

checkNeedsRedeployment ran the full workflow diff inside its REPEATABLE READ
transaction, holding a pooled connection idle in transaction during CPU work.
Reads (active version, normalized draft, variables, credential migrations)
stay in the snapshot; the pure diff now runs after commit.

* fix(mothership): build integration tool schemas once per process instead of per user every 5s

The integration catalog cached each viewer's full schema set for 5s, so nearly
every call rebuilt a user-tool schema for all ~5.3k exposed tools, sized the
~7.7MB result with JSON.stringify, structuredClone'd it, and re-tokenized
every tool for BM25 - all synchronously.

The schemas depend only on the registry and deployment config, so they are now
built once per (surface, personal mode), frozen, and shared; the build yields
between tools. Each request still applies the viewer's own block visibility and
permission policy, and the plan-dependent email tagline is read per user with
the same 5s TTL. Results are byte-identical. BM25 term counts are memoized per
shared schema object.

* refactor(db): tighten pool-isolation follow-ups

Assert each frozen-schema mutation throws, rank BM25 by query terms instead of
copying each tool's vocabulary, gate the reporting usage cache on
isPrimaryPoolClient like the auth caches, and trim comments.

* fix(api-key): read workspace key settings on the auth pool

Give the memory harness integration mock a dbFor so the billing pool resolves.

* test(db): close every pool the isolation suite opens

Author: waleedlatif1

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

@@ -47,7 +47,7 @@ vi.mock('@sim/db', async () => {
   })
   state.client = client
   const db = drizzle(client)
-  return { db, dbReplica: db }
+  return { db, dbReplica: db, dbFor: () => db }
 })
 
 /** Subscription and payment-provider fixtures; callback, ledger and replay code are real. */

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

@@ -47,7 +47,7 @@ vi.mock('@sim/db', async () => {
   })
   state.client = client
   const db = drizzle(client)
-  return { db, dbReplica: db }
+  return { db, dbReplica: db, dbFor: () => db }
 })
 
 /** Subscription and payment-provider fixtures; callback, ledger and replay code are real. */

apps/sim/executor/handlers/agent/memory-harness.integration.ts

@@ -30,8 +30,8 @@ const fixture = vi.hoisted(() => ({
   uploads: '',
 }))
 
-vi.mock('@sim/db', () => ({
-  db: new Proxy(
+vi.mock('@sim/db', () => {
+  const db = new Proxy(
     {},
     {
       get(_target, property) {
@@ -40,8 +40,9 @@ vi.mock('@sim/db', () => ({
         return typeof value === 'function' ? value.bind(fixture.database) : value
       },
     }
-  ),
-}))
+  )
+  return { db, dbFor: () => db }
+})
 vi.mock('@/lib/uploads/core/setup.server', () => ({
   get UPLOAD_DIR_SERVER() {
     return fixture.uploads

apps/sim/lib/api-key/service.ts

@@ -3,6 +3,7 @@ import { apiKey as apiKeyTable, user as userTable } from '@sim/db/schema'
 import { createLogger, setRequestAuth } from '@sim/logger'
 import { and, eq, isNull, lt, or } from 'drizzle-orm'
 import { hashApiKey } from '@/lib/api-key/crypto'
+import { getAuthPool } from '@/lib/auth/database-context'
 import { createDetachedTouch } from '@/lib/core/utils/background'
 import { getUserEntityPermissions } from '@/lib/workspaces/permissions/utils'
 import { getWorkspaceBillingSettings, type WorkspaceBillingSettings } from '@/lib/workspaces/utils'
@@ -55,14 +56,14 @@ export async function authenticateApiKeyFromHeader(
     let workspaceSettings: WorkspaceBillingSettings | null = null
 
     if (options.workspaceId) {
-      workspaceSettings = await getWorkspaceBillingSettings(options.workspaceId)
+      workspaceSettings = await getWorkspaceBillingSettings(options.workspaceId, getAuthPool())
       if (!workspaceSettings) {
         return { success: false, error: 'Workspace not found' }
       }
     }
 
     const keyHash = hashApiKey(apiKeyHeader)
-    const rows: HashCandidate[] = await db
+    const rows: HashCandidate[] = await getAuthPool()
       .select({
         id: apiKeyTable.id,
         userId: apiKeyTable.userId,

apps/sim/lib/auth/database-context.ts

@@ -1,12 +1,17 @@
 import { AsyncLocalStorage } from 'node:async_hooks'
-import { db } from '@sim/db'
+import { dbFor } from '@sim/db'
 import type { DbOrTx } from '@/lib/db/types'
 
 const authDatabase = new AsyncLocalStorage<DbOrTx>()
 
-/** The Drizzle executor backing the current Better Auth transaction, or the root database. */
+/** The auth pool's root client (see `DB_POOL_PROFILES.auth`). */
+export function getAuthPool() {
+  return dbFor('auth')
+}
+
+/** The Drizzle executor backing the current Better Auth transaction, or the auth pool. */
 export function getAuthDatabase(): DbOrTx {
-  return authDatabase.getStore() ?? db
+  return authDatabase.getStore() ?? getAuthPool()
 }
 
 /** Keeps database hooks on the adapter's connection so they can see uncommitted auth records. */

apps/sim/lib/auth/oauth-provider-adapter-guard.ts

@@ -1,10 +1,11 @@
 import { AsyncLocalStorage } from 'node:async_hooks'
-import { db } from '@sim/db'
+import type { db } from '@sim/db'
 import { oauthConsent, oauthTokenFamily } from '@sim/db/schema'
 import { generateId } from '@sim/utils/id'
 import type { drizzleAdapter } from 'better-auth/adapters/drizzle'
 import { APIError } from 'better-auth/api'
 import { inArray } from 'drizzle-orm'
+import { getAuthPool } from '@/lib/auth/database-context'
 import { getOAuthIssuedResource } from '@/lib/auth/oauth-resource'
 import { capabilityRefusal } from '@/lib/permission-groups/capabilities'
 import { isCapabilityWithheldForUser } from '@/lib/permission-groups/user-scope.server'
@@ -37,7 +38,7 @@ async function deleteTrackedFamilies(
  */
 export async function withOAuthProviderIssuanceCompensation<T extends Response>(
   work: () => Promise<T>,
-  database: AuthDatabase = db
+  database: AuthDatabase = getAuthPool()
 ): Promise<T> {
   const familyIds = new Set<string>()
   try {
@@ -87,7 +88,7 @@ function requireConsentInsert(data: Record<string, unknown>): OAuthConsentInsert
  */
 export function guardOAuthProviderWrites(
   adapter: BetterAuthAdapter,
-  database: AuthDatabase = db
+  database: AuthDatabase = getAuthPool()
 ): BetterAuthAdapter {
   return {
     ...adapter,

apps/sim/lib/auth/security-policy.test.ts

@@ -1,6 +1,6 @@
 import { db } from '@sim/db'
-import { dbChainMockFns, resetDbChainMock } from '@sim/testing'
-import { afterAll, beforeEach, describe, expect, it, vi } from 'vitest'
+import { databaseMockFns, dbChainMockFns, resetDbChainMock } from '@sim/testing'
+import { afterAll, afterEach, beforeEach, describe, expect, it, vi } from 'vitest'
 import {
   getMemberOrganizationId,
   getSecurityPolicyVersion,
@@ -146,3 +146,37 @@ describe('getSessionCookieCacheVersion', () => {
     expect(limit).toHaveBeenCalledTimes(2)
   })
 })
+
+describe('shared caches behind a sub-pool client', () => {
+  beforeEach(() => {
+    resetDbChainMock()
+    invalidateMembershipCache('pool-user')
+    invalidateSecurityPolicyVersionCache('pool-org')
+  })
+
+  afterEach(() => {
+    databaseMockFns.mockIsPrimaryPoolClient.mockReset()
+  })
+
+  /**
+   * Better Auth reads the cookie-cache version on every cached-session request through the auth
+   * pool, not `db`; caching only reads made through `db` would cost two queries per request.
+   */
+  it('serves a repeat cookie-cache version from cache when read through the auth pool', async () => {
+    const authPool: typeof db = Object.create(db)
+    databaseMockFns.mockIsPrimaryPoolClient.mockImplementation(
+      (executor) => executor === db || executor === authPool
+    )
+    dbChainMockFns.limit
+      .mockResolvedValueOnce([{ organizationId: 'pool-org' }])
+      .mockResolvedValueOnce([{ version: 3 }])
+
+    await expect(getSessionCookieCacheVersion({ userId: 'pool-user' }, authPool)).resolves.toBe(
+      'pool-org:3'
+    )
+    await expect(getSessionCookieCacheVersion({ userId: 'pool-user' }, authPool)).resolves.toBe(
+      'pool-org:3'
+    )
+    expect(dbChainMockFns.limit).toHaveBeenCalledTimes(2)
+  })
+})

apps/sim/lib/auth/security-policy.ts

@@ -1,4 +1,4 @@
-import { db } from '@sim/db'
+import { db, isPrimaryPoolClient } from '@sim/db'
 import { member, organization } from '@sim/db/schema'
 import { createLogger } from '@sim/logger'
 import { eq } from 'drizzle-orm'
@@ -50,7 +50,7 @@ export async function getSecurityPolicyVersion(
 ): Promise<number> {
   if (!organizationId) return DEFAULT_VERSION
 
-  const cached = executor === db ? versionCache.get(organizationId) : undefined
+  const cached = isPrimaryPoolClient(executor) ? versionCache.get(organizationId) : undefined
   if (cached !== undefined) return cached
 
   try {
@@ -61,7 +61,7 @@ export async function getSecurityPolicyVersion(
       .limit(1)
 
     const version = row?.version ?? DEFAULT_VERSION
-    if (executor === db) versionCache.set(organizationId, version)
+    if (isPrimaryPoolClient(executor)) versionCache.set(organizationId, version)
     return version
   } catch (error) {
     logger.error('Failed to resolve security policy version; using default', {
@@ -133,7 +133,7 @@ export async function getMemberOrganizationId(
   if (!userId) return null
 
   /** A transaction may have just created or moved the membership; shared cache entries cannot see it. */
-  const cached = executor === db ? membershipCache.get(userId) : undefined
+  const cached = isPrimaryPoolClient(executor) ? membershipCache.get(userId) : undefined
   if (cached) return cached.organizationId
 
   try {
@@ -144,7 +144,7 @@ export async function getMemberOrganizationId(
       .limit(1)
 
     const organizationId = row?.organizationId ?? null
-    if (executor === db) {
+    if (isPrimaryPoolClient(executor)) {
       membershipCache.set(userId, { organizationId }, { ttl: membershipCacheTtlMs(organizationId) })
     }
     return organizationId

apps/sim/lib/auth/session-policy.ts

@@ -1,4 +1,4 @@
-import { db } from '@sim/db'
+import { db, isPrimaryPoolClient } from '@sim/db'
 import type { SessionPolicySettings } from '@sim/db/schema'
 import { organization } from '@sim/db/schema'
 import { createLogger } from '@sim/logger'
@@ -57,7 +57,7 @@ export async function getSessionPolicy(
   if (!organizationId) return NO_POLICY
 
   /** Uncommitted policy reads must neither consume nor populate the shared cache. */
-  const cached = executor === db ? policyCache.get(organizationId) : undefined
+  const cached = isPrimaryPoolClient(executor) ? policyCache.get(organizationId) : undefined
   if (cached) return cached
 
   try {
@@ -78,7 +78,7 @@ export async function getSessionPolicy(
           idleTimeoutHours: settings.idleTimeoutHours ?? null,
         }
       : NO_POLICY
-    if (executor === db) policyCache.set(organizationId, policy)
+    if (isPrimaryPoolClient(executor)) policyCache.set(organizationId, policy)
     return policy
   } catch (error) {
     logger.error('Failed to resolve session policy; applying no policy', {

apps/sim/lib/auth/sim-auth-adapter.ts

@@ -1,10 +1,9 @@
-import { db } from '@sim/db'
 import * as schema from '@sim/db/schema'
 import type { BetterAuthOptions } from 'better-auth'
 import { drizzleAdapter } from 'better-auth/adapters/drizzle'
 import { APIError } from 'better-auth/api'
 import { and, eq } from 'drizzle-orm'
-import { runWithAuthDatabase } from '@/lib/auth/database-context'
+import { getAuthPool, runWithAuthDatabase } from '@/lib/auth/database-context'
 import {
   type AuthDatabase,
   guardOAuthProviderWrites,
@@ -66,7 +65,7 @@ function createTransactionAdapter(
  */
 export function createSimAuthAdapter(
   options: BetterAuthOptions,
-  database: AuthDatabase = db
+  database: AuthDatabase = getAuthPool()
 ): BetterAuthAdapter {
   const guarded = createGuardedAdapter(options, database)
   const transaction: BetterAuthAdapter['transaction'] = (callback) =>