diff --git a/apps/sim/app/api/billing/update-cost/replay.live.test.ts b/apps/sim/app/api/billing/update-cost/replay.live.test.ts index 44d50a6e642..2d7be627d8f 100644 --- a/apps/sim/app/api/billing/update-cost/replay.live.test.ts +++ b/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. */ diff --git a/apps/sim/app/api/billing/update-cost/worker.live.test.ts b/apps/sim/app/api/billing/update-cost/worker.live.test.ts index 9729f86ffa1..d62914d9b86 100644 --- a/apps/sim/app/api/billing/update-cost/worker.live.test.ts +++ b/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. */ diff --git a/apps/sim/executor/handlers/agent/memory-harness.integration.ts b/apps/sim/executor/handlers/agent/memory-harness.integration.ts index af2b350d8a7..c6fd1683db0 100644 --- a/apps/sim/executor/handlers/agent/memory-harness.integration.ts +++ b/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 diff --git a/apps/sim/lib/api-key/service.ts b/apps/sim/lib/api-key/service.ts index d8edaa27d00..8dbe94cf520 100644 --- a/apps/sim/lib/api-key/service.ts +++ b/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, diff --git a/apps/sim/lib/auth/database-context.ts b/apps/sim/lib/auth/database-context.ts index 5c1326dd004..bb55187da3d 100644 --- a/apps/sim/lib/auth/database-context.ts +++ b/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() -/** 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. */ diff --git a/apps/sim/lib/auth/oauth-provider-adapter-guard.ts b/apps/sim/lib/auth/oauth-provider-adapter-guard.ts index 72bae75102c..2ca707e0e34 100644 --- a/apps/sim/lib/auth/oauth-provider-adapter-guard.ts +++ b/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( work: () => Promise, - database: AuthDatabase = db + database: AuthDatabase = getAuthPool() ): Promise { const familyIds = new Set() try { @@ -87,7 +88,7 @@ function requireConsentInsert(data: Record): OAuthConsentInsert */ export function guardOAuthProviderWrites( adapter: BetterAuthAdapter, - database: AuthDatabase = db + database: AuthDatabase = getAuthPool() ): BetterAuthAdapter { return { ...adapter, diff --git a/apps/sim/lib/auth/security-policy.test.ts b/apps/sim/lib/auth/security-policy.test.ts index be96063c022..d4b97c6ec55 100644 --- a/apps/sim/lib/auth/security-policy.test.ts +++ b/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) + }) +}) diff --git a/apps/sim/lib/auth/security-policy.ts b/apps/sim/lib/auth/security-policy.ts index 354bd531901..8b8e455e289 100644 --- a/apps/sim/lib/auth/security-policy.ts +++ b/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 { 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 diff --git a/apps/sim/lib/auth/session-policy.ts b/apps/sim/lib/auth/session-policy.ts index 8508d9caa6f..0f28b9e53e5 100644 --- a/apps/sim/lib/auth/session-policy.ts +++ b/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', { diff --git a/apps/sim/lib/auth/sim-auth-adapter.ts b/apps/sim/lib/auth/sim-auth-adapter.ts index b2bebd08c7c..741a211debd 100644 --- a/apps/sim/lib/auth/sim-auth-adapter.ts +++ b/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) => diff --git a/apps/sim/lib/auth/sso-policy.ts b/apps/sim/lib/auth/sso-policy.ts index d2db3ace8d8..0fa8c32d267 100644 --- a/apps/sim/lib/auth/sso-policy.ts +++ b/apps/sim/lib/auth/sso-policy.ts @@ -1,4 +1,4 @@ -import { db } from '@sim/db' +import { db, isPrimaryPoolClient } from '@sim/db' import { organization } from '@sim/db/schema' import { createLogger } from '@sim/logger' import { APIError } from 'better-auth/api' @@ -43,7 +43,7 @@ export async function isSsoRequiredForOrganization( if (!organizationId) return false /** Uncommitted reads must neither consume nor populate the shared cache. */ - const cached = executor === db ? requirementCache.get(organizationId) : undefined + const cached = isPrimaryPoolClient(executor) ? requirementCache.get(organizationId) : undefined if (cached !== undefined) return cached const [row] = await executor @@ -64,7 +64,7 @@ export async function isSsoRequiredForOrganization( onError: 'throw', })) && (await hasSignInCapableSsoProvider(organizationId, executor)) - if (executor === db) requirementCache.set(organizationId, required) + if (isPrimaryPoolClient(executor)) requirementCache.set(organizationId, required) return required } diff --git a/apps/sim/lib/billing/core/reporting-usage-cache.integration.ts b/apps/sim/lib/billing/core/reporting-usage-cache.integration.ts index 062c9de0f98..c34b1f761a9 100644 --- a/apps/sim/lib/billing/core/reporting-usage-cache.integration.ts +++ b/apps/sim/lib/billing/core/reporting-usage-cache.integration.ts @@ -23,7 +23,15 @@ const { transaction, summedPeriods } = vi.hoisted(() => ({ const databaseUrl = readTestDatabaseUrl() const redisUrl = readTestRedisUrl() -vi.mock('@sim/db', () => ({ db: { transaction }, dbReplica: {} })) +vi.mock('@sim/db', () => { + const db = { transaction } + return { + db, + dbFor: () => ({ transaction }), + dbReplica: {}, + isPrimaryPoolClient: (executor: object) => executor === db, + } +}) vi.mock('@/lib/core/config/redis', () => redisConfigMock) vi.mock('@/lib/billing/core/usage-log', async (importOriginal) => { const actual = await importOriginal() diff --git a/apps/sim/lib/billing/core/reporting-usage-cache.ts b/apps/sim/lib/billing/core/reporting-usage-cache.ts index 60dc97d23fe..507cb78066f 100644 --- a/apps/sim/lib/billing/core/reporting-usage-cache.ts +++ b/apps/sim/lib/billing/core/reporting-usage-cache.ts @@ -1,4 +1,4 @@ -import { db } from '@sim/db' +import { db, isPrimaryPoolClient } from '@sim/db' import { createLogger } from '@sim/logger' import { getErrorMessage } from '@sim/utils/errors' import { generateShortId } from '@sim/utils/id' @@ -417,7 +417,7 @@ export function readSoftGateUsageCost( period: UsageQueryPeriod, executor: DbClient = db ): Promise { - if (isReportingPeriod(period) && executor === db) { + if (isReportingPeriod(period) && isPrimaryPoolClient(executor)) { return readCachedReportingUsageCost(entity, period) } return getBillingPeriodUsageCost(entity, period, undefined, executor) diff --git a/apps/sim/lib/billing/core/usage-log.integration.ts b/apps/sim/lib/billing/core/usage-log.integration.ts index b6866e49665..4f91dad0d5b 100644 --- a/apps/sim/lib/billing/core/usage-log.integration.ts +++ b/apps/sim/lib/billing/core/usage-log.integration.ts @@ -19,7 +19,7 @@ import { afterAll, beforeAll, beforeEach, describe, expect, it, vi } from 'vites const { transaction } = vi.hoisted(() => ({ transaction: vi.fn() })) const databaseUrl = readTestDatabaseUrl() -vi.mock('@sim/db', () => ({ db: { transaction }, dbReplica: {} })) +vi.mock('@sim/db', () => ({ db: { transaction }, dbFor: () => ({ transaction }), dbReplica: {} })) vi.mock('@/lib/billing/core/plan', () => ({ getHighestPrioritySubscription: vi.fn() })) vi.mock('@/lib/billing/subscriptions/utils', async (importOriginal) => ({ ...(await importOriginal()), diff --git a/apps/sim/lib/billing/core/usage-log.ts b/apps/sim/lib/billing/core/usage-log.ts index 2c29a25e4f1..837ab470c00 100644 --- a/apps/sim/lib/billing/core/usage-log.ts +++ b/apps/sim/lib/billing/core/usage-log.ts @@ -1,5 +1,5 @@ import { createHash } from 'node:crypto' -import { db, dbReplica } from '@sim/db' +import { db, dbFor, dbReplica } from '@sim/db' import { subscription as subscriptionTable, usageLog, workflow } from '@sim/db/schema' import { createLogger } from '@sim/logger' import { toNumberOrNull } from '@sim/utils/coerce' @@ -475,7 +475,7 @@ export async function recordUsage(params: RecordUsageParams): Promise { const context = await resolveBillingContext(userId, billingEntity, billingPeriod) - const insertedRows = await (tx ?? db) + const insertedRows = await (tx ?? dbFor('billing')) .insert(usageLog) .values( validEntries.map((entry, index) => { @@ -842,7 +842,7 @@ export async function recordCumulativeUsage( } try { - const result = await db.transaction(async (tx) => { + const result = await dbFor('billing').transaction(async (tx) => { enterStage('configure') await tx.execute(sql` select diff --git a/apps/sim/lib/billing/threshold-billing.ts b/apps/sim/lib/billing/threshold-billing.ts index 1f50d5b2050..afb72fb6a36 100644 --- a/apps/sim/lib/billing/threshold-billing.ts +++ b/apps/sim/lib/billing/threshold-billing.ts @@ -1,5 +1,5 @@ import { AuditAction, AuditResourceType, recordAudit } from '@sim/audit' -import { db } from '@sim/db' +import { db, dbFor } from '@sim/db' import { member, organization, subscription, userStats } from '@sim/db/schema' import { createLogger } from '@sim/logger' import { toError } from '@sim/utils/errors' @@ -343,7 +343,7 @@ export async function checkAndBillOverageThreshold( const billingPeriod = new Date(periodEnd * 1000).toISOString().slice(0, 7) const totalOverageCents = Math.round(currentOverage * 100) - const billedResult = await db.transaction( + const billedResult = await dbFor('billing').transaction( async (tx): Promise => { await tx.execute(sql.raw(`SET LOCAL lock_timeout = '${BILLING_LOCK_TIMEOUT_MS}ms'`)) @@ -678,7 +678,7 @@ async function checkAndBillOrganizationOverageThreshold( const billingPeriod = new Date(periodEnd * 1000).toISOString().slice(0, 7) const totalOverageCents = Math.round(currentOverage * 100) - const orgBilledResult = await db.transaction( + const orgBilledResult = await dbFor('billing').transaction( async ( tx ): Promise< diff --git a/apps/sim/lib/core/config/env.ts b/apps/sim/lib/core/config/env.ts index 423c949174b..0157cf43222 100644 --- a/apps/sim/lib/core/config/env.ts +++ b/apps/sim/lib/core/config/env.ts @@ -121,6 +121,8 @@ export const env = createEnv({ DATABASE_URL_REALTIME: z.string().url().optional(), // Per-role primary URL override (realtime) DATABASE_URL_CLEANUP: z.string().url().optional(), // Sub-process pool URL override (cleanup jobs, via dbFor) DATABASE_URL_EXEC: z.string().url().optional(), // Sub-process pool URL override (inline execution writes, via dbFor) + DATABASE_URL_BILLING: z.string().url().optional(), // Sub-process pool URL override (usage-ledger writes, via dbFor) + DATABASE_URL_AUTH: z.string().url().optional(), // Sub-process pool URL override (session and API-key lookups, via dbFor) DATABASE_REPLICA_URL_WEB: z.string().url().optional(), // Per-role replica URL override; falls back to DATABASE_REPLICA_URL DATABASE_REPLICA_URL_TRIGGER: z.string().url().optional(), // Per-role replica URL override (trigger) DATABASE_REPLICA_URL_REALTIME: z.string().url().optional(), // Per-role replica URL override (realtime) diff --git a/apps/sim/lib/db/pool-isolation.integration.ts b/apps/sim/lib/db/pool-isolation.integration.ts new file mode 100644 index 00000000000..d21c48fc952 --- /dev/null +++ b/apps/sim/lib/db/pool-isolation.integration.ts @@ -0,0 +1,145 @@ +/** + * Real `@sim/db` pools against the integration database: with every connection of the process + * pool stuck behind a held advisory lock, usage-ledger writes and auth lookups must still finish + * on their own pools. postgres.js pipelines new work onto busy connections when none is idle, so + * anything still on the process pool waits for the lock instead. + */ +import { db, dbFor } from '@sim/db' +import { session, user } from '@sim/db/schema' +import { readTestDatabaseUrl } from '@sim/db/testing/test-infrastructure' +import { sleep } from '@sim/utils/helpers' +import { generateId } from '@sim/utils/id' +import { randomInt } from '@sim/utils/random' +import { eq, sql } from 'drizzle-orm' +import postgres from 'postgres' +import { afterAll, beforeAll, describe, expect, it } from 'vitest' +import { authenticateApiKeyFromHeader } from '@/lib/api-key/service' +import { createSimAuthAdapter } from '@/lib/auth/sim-auth-adapter' +import { recordCumulativeUsage } from '@/lib/billing/core/usage-log' + +/** Covers a new pool's first connect; a stalled operation waits for the lock, which never comes. */ +const ISOLATED_BUDGET_MS = 1_500 +const SATURATION_PROBE_MS = 300 + +const control = postgres(readTestDatabaseUrl(), { max: 1, onnotice: () => {} }) +const lockKey = randomInt(1, 2 ** 31) +const userId = `pool-isolation-${generateId()}` +const sessionToken = `pool-isolation-${generateId()}` +const now = new Date() + +let saturation: Promise | undefined + +async function waitForSaturatedProcessPool(connections: number): Promise { + for (let attempt = 0; attempt < 400; attempt++) { + const [{ waiting }] = await control<{ waiting: number }[]>` + SELECT count(*)::int AS waiting FROM pg_stat_activity + WHERE datname = current_database() AND wait_event = 'advisory' + AND query LIKE '%pg_advisory_xact_lock_shared%'` + if (waiting >= connections) return + await sleep(5) + } + throw new Error('The process pool never filled with lock waiters') +} + +/** Resolves `'done'` when `work` settles within the budget, `'stalled'` otherwise. */ +function settlesWithin(work: Promise, budgetMs: number): Promise<'done' | 'stalled'> { + return Promise.race([ + work.then(() => 'done' as const), + sleep(budgetMs).then(() => 'stalled' as const), + ]) +} + +beforeAll(async () => { + await db.insert(user).values({ + id: userId, + name: 'Pool isolation', + email: `${userId}@example.com`, + emailVerified: true, + createdAt: now, + updatedAt: now, + }) + await db.insert(session).values({ + id: generateId(), + token: sessionToken, + userId, + expiresAt: new Date(now.getTime() + 60 * 60 * 1000), + createdAt: now, + updatedAt: now, + }) + + await control`SELECT pg_advisory_lock(${lockKey})` + const connections = db.$client.options.max + saturation = Promise.all( + Array.from({ length: connections }, () => + db.execute(sql`SELECT pg_advisory_xact_lock_shared(${lockKey})`) + ) + ) + await waitForSaturatedProcessPool(connections) +}) + +afterAll(async () => { + try { + await control`SELECT pg_advisory_unlock(${lockKey})` + await saturation + await db.delete(user).where(eq(user.id, userId)) + } finally { + await Promise.all([ + db.$client.end(), + dbFor('billing').$client.end(), + dbFor('auth').$client.end(), + control.end(), + ]) + } +}) + +describe('dedicated pools under a saturated process pool', () => { + it('leaves the process pool unable to serve a new query', async () => { + await expect(settlesWithin(db.execute(sql`SELECT 1`), SATURATION_PROBE_MS)).resolves.toBe( + 'stalled' + ) + }) + + it('records cumulative usage on the billing pool', async () => { + const write = recordCumulativeUsage({ + userId, + billingEntity: { type: 'user', id: userId }, + billingPeriod: { + start: new Date('2026-10-01T00:00:00Z'), + end: new Date('2026-11-01T00:00:00Z'), + }, + source: 'workflow', + model: 'pool-isolation-model', + cost: 0.25, + eventKey: `pool-isolation:${generateId()}`, + }) + + await expect(settlesWithin(write, ISOLATED_BUDGET_MS)).resolves.toBe('done') + await expect(write).resolves.toMatchObject({ billed: true, total: 0.25 }) + }) + + it('reads a session through the auth adapter on the auth pool', async () => { + const lookup = createSimAuthAdapter({}).findOne<{ userId: string }>({ + model: 'session', + where: [{ field: 'token', value: sessionToken }], + }) + + await expect(settlesWithin(lookup, ISOLATED_BUDGET_MS)).resolves.toBe('done') + await expect(lookup).resolves.toMatchObject({ userId }) + }) + + it('looks up an API key on the auth pool', async () => { + const lookup = authenticateApiKeyFromHeader(`sk-sim-${generateId()}`) + + await expect(settlesWithin(lookup, ISOLATED_BUDGET_MS)).resolves.toBe('done') + await expect(lookup).resolves.toEqual({ success: false, error: 'Invalid API key' }) + }) + + it('reads workspace key settings on the auth pool', async () => { + const lookup = authenticateApiKeyFromHeader(`sk-sim-${generateId()}`, { + workspaceId: generateId(), + }) + + await expect(settlesWithin(lookup, ISOLATED_BUDGET_MS)).resolves.toBe('done') + await expect(lookup).resolves.toEqual({ success: false, error: 'Workspace not found' }) + }) +}) diff --git a/apps/sim/lib/integrations/tool-projection.ts b/apps/sim/lib/integrations/tool-projection.ts index 085d4ce0608..a7fa47c02da 100644 --- a/apps/sim/lib/integrations/tool-projection.ts +++ b/apps/sim/lib/integrations/tool-projection.ts @@ -76,22 +76,6 @@ export function projectIntegrationToolsForViewer( return { tools, allowedBlockTypes, isToolAllowed } } -/** - * Stable signature of the policy {@link projectIntegrationToolsForViewer} reads, - * for keying caches whose contents depend on the projection. - * - * The projection is only as fresh as what keys it: an entry cached under a - * viewer's identity alone outlives the policy that produced it, so an admin's - * change would not take effect until the entry expired. Mirrors - * `visibilitySignature`, which does the same job for block visibility. - */ -export function integrationGateSignature(config: IntegrationGateConfig | null | undefined): string { - return JSON.stringify([ - config?.allowedIntegrations ? [...config.allowedIntegrations].sort() : null, - config?.deniedTools?.length ? [...config.deniedTools].sort() : null, - ]) -} - /** What a viewer's `deniedTools` denylist costs the block schemas they are shown. */ export interface DeniedBlockOperations { /** diff --git a/apps/sim/lib/logs/execution/logger.ts b/apps/sim/lib/logs/execution/logger.ts index 0c26186d937..e1bdc8e0e10 100644 --- a/apps/sim/lib/logs/execution/logger.ts +++ b/apps/sim/lib/logs/execution/logger.ts @@ -1773,7 +1773,7 @@ export class ExecutionLogger { // flow and only matters under a cross-process double-completion of the // same execution, where it stops a stale already-billed read from // dropping the larger delta. - await db.transaction(async (tx) => { + await dbFor('billing').transaction(async (tx) => { await tx.execute( sql`select set_config('lock_timeout', ${`${USAGE_RECONCILE_LOCK_TIMEOUT_MS}ms`}, true)` ) diff --git a/apps/sim/lib/mothership/block-visibility.ts b/apps/sim/lib/mothership/block-visibility.ts index f35e787de45..430fe099fe9 100644 --- a/apps/sim/lib/mothership/block-visibility.ts +++ b/apps/sim/lib/mothership/block-visibility.ts @@ -46,15 +46,3 @@ export function getBlockVisibilityForCopilot( } return promise } - -/** - * Stable signature of a visibility state, for keying caches whose contents - * depend on the gated projection (e.g. the integration tool-schema LRU). - */ -export function visibilitySignature(vis: BlockVisibilityState): string { - return JSON.stringify([ - [...vis.revealed].sort(), - [...vis.disabled].sort(), - [...vis.previewTagged].sort(), - ]) -} diff --git a/apps/sim/lib/mothership/chat/payload.test.ts b/apps/sim/lib/mothership/chat/payload.test.ts index 61fef4c6fcc..d4c5107a772 100644 --- a/apps/sim/lib/mothership/chat/payload.test.ts +++ b/apps/sim/lib/mothership/chat/payload.test.ts @@ -22,7 +22,7 @@ import { workspaceFileManagerMock, workspaceFileManagerMockFns, } from '@sim/testing/mocks/workspace-file-manager.mock' -import { beforeEach, describe, expect, it, vi } from 'vitest' +import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest' import { getExposedIntegrationTools } from '@/lib/integrations/tool-catalog' import { ChatPayloadSchema, TOOL_DEADLINE_CAPABILITY } from '@/lib/mothership/generated/protocol' import { searchUsersV2Tool } from '@/tools/github/search_users' @@ -128,7 +128,6 @@ vi.mock('@/lib/mothership/block-visibility', () => ({ disabled: new Set(), previewTagged: new Set(), })), - visibilitySignature: vi.fn(() => 'vis:none'), })) vi.mock('@/lib/integrations/tool-catalog', () => ({ @@ -273,38 +272,85 @@ describe('buildIntegrationToolSchemas', () => { expect(mockCreateUserToolSchema).not.toHaveBeenCalled() }) - it('briefly reuses built schemas for the same user and surface', async () => { + it('builds each tool schema once per process, not once per viewer', async () => { mockGetHighestPrioritySubscription.mockResolvedValue({ plan: 'pro', status: 'active' }) + mockGetUserPermissionConfig.mockResolvedValueOnce({ allowedIntegrations: ['gmail'] }) - const first = await buildIntegrationToolSchemas('user-cache') - first[0].input_schema.mutated = true - if (first[0].outputs) first[0].outputs.mutated = true - const second = await buildIntegrationToolSchemas('user-cache') + await buildIntegrationToolSchemas('user-a', { schemaSurface: 'copilot' }, 'workspace-a') + await buildIntegrationToolSchemas('user-b', { schemaSurface: 'copilot' }, 'workspace-b') + await buildIntegrationToolSchemas('user-b', { schemaSurface: 'copilot' }, 'workspace-b') - expect(mockGetHighestPrioritySubscription).toHaveBeenCalledTimes(1) expect(mockCreateUserToolSchema).toHaveBeenCalledTimes(3) - expect(second[0].input_schema).not.toHaveProperty('mutated') - expect(second[0].outputs).not.toHaveProperty('mutated') }) - it('isolates nested schemas and required fields between requests', async () => { - mockCreateUserToolSchema.mockReturnValueOnce({ + it('keeps one caller from changing what any other viewer receives', async () => { + mockGetHighestPrioritySubscription.mockResolvedValue({ plan: 'pro', status: 'active' }) + mockCreateUserToolSchema.mockReturnValue({ type: 'object', properties: { recipients: { type: 'array', items: { type: 'string' } } }, required: ['recipients'], }) - const first = await buildIntegrationToolSchemas('user-nested-schema') - const properties = first[0].input_schema.properties as Record - properties.recipients = { type: 'number' } - ;(first[0].input_schema.required as string[]).push('forged') - - const second = await buildIntegrationToolSchemas('user-nested-schema') + const first = await buildIntegrationToolSchemas('user-mutating') + first.push({ name: 'mcp-private-tool', description: '', input_schema: {} }) + expect(() => { + ;(first[0].input_schema.properties as Record).recipients = { type: 'number' } + }).toThrow(TypeError) + expect(() => (first[0].input_schema.required as string[]).push('forged')).toThrow(TypeError) + expect(() => { + first[0].description = 'forged' + }).toThrow(TypeError) + + const second = await buildIntegrationToolSchemas('user-reading') + expect(second.map((tool) => tool.name)).toEqual([ + 'gmail_send', + 'brandfetch_search', + 'run_workflow', + ]) + expect(second[0].description).not.toBe('forged') expect(second[0].input_schema.properties).toEqual({ recipients: { type: 'array', items: { type: 'string' } }, }) expect(second[0].input_schema.required).toEqual(['recipients']) }) + it('appends the free-plan email tagline only for the viewer whose plan earns it', async () => { + mockGetHighestPrioritySubscription + .mockResolvedValueOnce(null) + .mockResolvedValueOnce({ plan: 'pro', status: 'active' }) + const tagline = 'sent with sim ai' + const describeGmail = (tools: Array<{ name: string; description: string }>) => + tools.find((tool) => tool.name === 'gmail_send')?.description ?? '' + + const free = await buildIntegrationToolSchemas('free-user') + const paid = await buildIntegrationToolSchemas('paid-user') + const freeAgain = await buildIntegrationToolSchemas('free-user') + + expect(describeGmail(free)).toContain(tagline) + expect(describeGmail(paid)).not.toContain(tagline) + expect(describeGmail(freeAgain)).toContain(tagline) + }) + + it("gates the shared schemas by each viewer's own permission policy", async () => { + mockGetHighestPrioritySubscription.mockResolvedValue({ plan: 'pro', status: 'active' }) + mockGetUserPermissionConfig + .mockResolvedValueOnce({ allowedIntegrations: null, deniedTools: ['gmail_send'] }) + .mockResolvedValueOnce(null) + + const restricted = await buildIntegrationToolSchemas( + 'restricted-user', + { schemaSurface: 'copilot' }, + 'workspace-shared' + ) + const unrestricted = await buildIntegrationToolSchemas( + 'open-user', + { schemaSurface: 'copilot' }, + 'workspace-shared' + ) + + expect(restricted.map((tool) => tool.name)).not.toContain('gmail_send') + expect(unrestricted.map((tool) => tool.name)).toContain('gmail_send') + }) + it('coalesces simultaneous catalog builds', async () => { const catalogs = await Promise.all( Array.from({ length: 20 }, () => buildIntegrationToolSchemas('concurrent-user')) @@ -323,7 +369,6 @@ describe('buildIntegrationToolSchemas', () => { ) const recovered = await buildIntegrationToolSchemas('schema-failure') expect(recovered).toHaveLength(3) - expect(mockGetHighestPrioritySubscription).toHaveBeenCalledTimes(2) }) it('propagates subscription failures without caching a guessed catalog', async () => { @@ -335,18 +380,6 @@ describe('buildIntegrationToolSchemas', () => { expect(await buildIntegrationToolSchemas('subscription-failure')).toHaveLength(3) }) - it('evicts catalogs by their byte size before reaching the entry cap', async () => { - mockCreateUserToolSchema.mockImplementation(() => ({ - type: 'object', - properties: { large: { type: 'string', description: 'x'.repeat(1024 * 1024) } }, - })) - for (let i = 0; i < 12; i++) await buildIntegrationToolSchemas(`sized-user-${i}`) - expect(mockGetHighestPrioritySubscription).toHaveBeenCalledTimes(12) - await buildIntegrationToolSchemas('sized-user-0') - expect(mockGetHighestPrioritySubscription).toHaveBeenCalledTimes(13) - mockCreateUserToolSchema.mockImplementation(() => ({ type: 'object', properties: {} })) - }) - it('rebuilds instead of serving a cache entry from the previous policy', async () => { mockGetHighestPrioritySubscription.mockResolvedValue({ plan: 'pro', status: 'active' }) mockGetUserPermissionConfig.mockResolvedValue({ allowedIntegrations: null, deniedTools: [] }) @@ -660,6 +693,7 @@ describe('buildCopilotRequestPayload', () => { }) describe('Assistant payload', () => { + afterEach(() => vi.mocked(getExposedIntegrationTools).mockReset()) beforeEach(() => { resetEnvFlagsMock() mockGetUserPermissionConfig.mockResolvedValue(null) @@ -671,7 +705,7 @@ describe('Assistant payload', () => { it('discovers the existing GitHub user lookup tool with a personal credential in live Search', async () => { clearIntegrationToolSchemaCacheForTests() mockSearchApprovals.mockResolvedValue(new Map([['github', true]])) - vi.mocked(getExposedIntegrationTools).mockReturnValueOnce([ + vi.mocked(getExposedIntegrationTools).mockReturnValue([ { toolId: searchUsersV2Tool.id, config: searchUsersV2Tool, @@ -706,7 +740,7 @@ describe('Assistant payload', () => { }) it('advertises approved personal organization integrations and rechecks revocation', async () => { clearIntegrationToolSchemaCacheForTests() - vi.mocked(getExposedIntegrationTools).mockReturnValueOnce([ + vi.mocked(getExposedIntegrationTools).mockReturnValue([ { toolId: gmailListLabelsV2Tool.id, config: gmailListLabelsV2Tool, diff --git a/apps/sim/lib/mothership/chat/payload.ts b/apps/sim/lib/mothership/chat/payload.ts index df4b74aa1f1..85fdef2fd74 100644 --- a/apps/sim/lib/mothership/chat/payload.ts +++ b/apps/sim/lib/mothership/chat/payload.ts @@ -7,12 +7,12 @@ import { LRUCache } from 'lru-cache' import { getHighestPrioritySubscription } from '@/lib/billing/core/subscription' import { isPaid } from '@/lib/billing/plan-helpers' import { isComputerUseAvailableFor } from '@/lib/computer-use/application/availability' -import type { BlockVisibilityState } from '@/lib/core/config/block-visibility' import { isHosted } from '@/lib/core/config/env-flags' +import { yieldWhenSliceExpires } from '@/lib/core/utils/event-loop' import { isOAuthServiceDeploymentAvailable } from '@/lib/integrations/availability.server' +import { getExposedIntegrationTools } from '@/lib/integrations/tool-catalog' import { type IntegrationGateConfig, - integrationGateSignature, projectIntegrationToolsForViewer, } from '@/lib/integrations/tool-projection' import type { WorkspaceSearchFilters } from '@/lib/knowledge/search/filters' @@ -23,10 +23,7 @@ import { isAssistantIntegrationParameter, isAssistantIntegrationTool, } from '@/lib/mothership/assistant/tool-policy' -import { - getBlockVisibilityForCopilot, - visibilitySignature, -} from '@/lib/mothership/block-visibility' +import { getBlockVisibilityForCopilot } from '@/lib/mothership/block-visibility' import type { AssistantImageContent } from '@/lib/mothership/chat/assistant-images' import { buildUploadedFileContext } from '@/lib/mothership/chat/upload-context' import { buildWorkspaceInventory } from '@/lib/mothership/chat/workspace-inventory' @@ -51,9 +48,10 @@ import { deriveHostedApiKeySupport } from '@/tools/hosted-api-key' import { getToolMetadata } from '@/tools/metadata' const logger = createLogger('CopilotChatPayload') -const INTEGRATION_TOOL_SCHEMA_CACHE_TTL_MS = 5_000 -const INTEGRATION_TOOL_SCHEMA_CACHE_MAX_ENTRIES = 500 -const INTEGRATION_TOOL_SCHEMA_CACHE_MAX_BYTES = 32 * 1024 * 1024 +const EMAIL_TAGLINE_CACHE_TTL_MS = 5_000 +const EMAIL_TAGLINE_CACHE_MAX_ENTRIES = 20_000 +/** One per schema surface and personal-accounts mode; the cap only backstops a new surface. */ +const INTEGRATION_TOOL_SCHEMA_VARIANT_MAX_ENTRIES = 8 interface BuildPayloadParams { message: string @@ -133,43 +131,56 @@ interface BuildIntegrationToolSchemasOptions { organizationId?: string } -interface IntegrationToolSchemaBuildContext { - userId: string - options: Required> - vis: BlockVisibilityState | null - permissionConfig: IntegrationGateConfig | null +type IntegrationToolSchemaVariantOptions = Required< + Omit +> + +/** + * Every exposed integration tool's schema for one surface and mode, by tool id. + * + * Depends only on the tool registry and deployment config (`isHosted`, OAuth availability), + * fixed for the process's life, never on the viewer, so it is built once per process and + * shared: the entries are frozen and must not be mutated. + * Viewer gating (block visibility, permission groups, the deployment allowlist) picks + * which entries a request returns. Tools whose description gains the free-plan email + * tagline also get a tagged copy in `withEmailTagline`, sharing the same input schema. + */ +interface IntegrationToolSchemaVariant { + schemas: ReadonlyMap + withEmailTagline: ReadonlyMap } -const integrationToolSchemaCache = new LRUCache< +const integrationToolSchemaVariants = new LRUCache< string, - ToolSchema[], - IntegrationToolSchemaBuildContext + IntegrationToolSchemaVariant, + IntegrationToolSchemaVariantOptions >({ - max: INTEGRATION_TOOL_SCHEMA_CACHE_MAX_ENTRIES, - maxSize: INTEGRATION_TOOL_SCHEMA_CACHE_MAX_BYTES, - sizeCalculation: (schemas) => Buffer.byteLength(JSON.stringify(schemas)), - ttl: INTEGRATION_TOOL_SCHEMA_CACHE_TTL_MS, + max: INTEGRATION_TOOL_SCHEMA_VARIANT_MAX_ENTRIES, ignoreFetchAbort: true, - fetchMethod: async (_key, _staleValue, { context }) => - buildIntegrationToolSchemasUncached(context), + fetchMethod: async (_key, _staleValue, { context }) => buildIntegrationToolSchemaVariant(context), }) -function getIntegrationToolSchemaCacheKey( - userId: string, - workspaceId: string | undefined, - schemaSurface: string, - visSignature: string, - gateSignature: string -): string { - // The visibility signature keys the entry to the viewer's gated projection — - // two users in one workspace with different preview reveals must not share. - // The gate signature does the same for permission-group policy, so an admin's - // change takes effect on the next build rather than when the entry expires. - return JSON.stringify([userId, workspaceId ?? null, schemaSurface, visSignature, gateSignature]) -} +/** Whether the viewer's plan adds the email tagline; as fresh as the subscription read it caches. */ +const emailTaglineCache = new LRUCache({ + max: EMAIL_TAGLINE_CACHE_MAX_ENTRIES, + ttl: EMAIL_TAGLINE_CACHE_TTL_MS, + ignoreFetchAbort: true, + fetchMethod: async (userId) => { + const subscription = await getHighestPrioritySubscription(userId) + return !subscription || !isPaid(subscription.plan) + }, +}) export function clearIntegrationToolSchemaCacheForTests(): void { - integrationToolSchemaCache.clear() + integrationToolSchemaVariants.clear() + emailTaglineCache.clear() +} + +function deepFreeze(value: T): T { + if (value === null || typeof value !== 'object' || Object.isFrozen(value)) return value + for (const child of Object.values(value)) deepFreeze(child) + Object.freeze(value) + return value } /** @@ -180,6 +191,9 @@ export function clearIntegrationToolSchemaCacheForTests(): void { * When `workspaceId` is provided the user's workspace permission config is * loaded once and used to skip any tool whose owning block is not in the * workspace's `allowedIntegrations` allowlist. + * + * The returned array is the caller's own, but its schemas are shared across + * requests and frozen: copy one before changing it. */ export async function buildIntegrationToolSchemas( userId: string, @@ -189,9 +203,6 @@ export async function buildIntegrationToolSchemas( const schemaSurface = options.schemaSurface ?? 'copilot' const personalAccountsOnly = options.personalAccountsOnly ?? false const vis = await getBlockVisibilityForCopilot(userId, workspaceId, options.organizationId) - // Resolved before the key, not inside the cached build, so the entry is keyed - // to the policy it was produced under. The read this adds is cheap next to - // what the entry caches: a user-tool schema per exposed integration tool. let permissionConfig: IntegrationGateConfig | null = null if (workspaceId) { const { resolvePermissionGroupConfig } = await import( @@ -204,17 +215,25 @@ export async function buildIntegrationToolSchemas( return [] permissionConfig = organizationConfig } - const cacheKey = getIntegrationToolSchemaCacheKey( - userId, - workspaceId, - personalAccountsOnly ? `${schemaSurface}:personal` : schemaSurface, - visibilitySignature(vis), - integrationGateSignature(permissionConfig) - ) - const schemas = await integrationToolSchemaCache.fetch(cacheKey, { - context: { userId, options: { schemaSurface, personalAccountsOnly }, vis, permissionConfig }, - }) - if (!schemas) throw new Error('Integration tool catalog is unavailable') + const [variant, appendEmailTagline] = await Promise.all([ + integrationToolSchemaVariants.fetch(JSON.stringify([schemaSurface, personalAccountsOnly]), { + context: { schemaSurface, personalAccountsOnly }, + }), + emailTaglineCache.fetch(userId), + ]) + if (!variant || appendEmailTagline === undefined) + throw new Error('Integration tool catalog is unavailable') + + const schemas: ToolSchema[] = [] + for (const { toolId, service } of projectIntegrationToolsForViewer(vis, permissionConfig).tools) { + const schema = + (appendEmailTagline ? variant.withEmailTagline.get(toolId) : undefined) ?? + variant.schemas.get(toolId) + if (!schema) continue + // The viewer's visible owning block can be a different service than the canonical one. + schemas.push(schema.service === service ? schema : Object.freeze({ ...schema, service })) + } + if (options.organizationId && personalAccountsOnly) { const approvals = await listOrganizationSearchApprovals(options.organizationId) const providers = new Set( @@ -222,36 +241,30 @@ export async function buildIntegrationToolSchemas( (connector) => connector.providerId ) ) - return structuredClone( - schemas.filter((schema) => { - const original = getToolMetadata(schema.name) - const metadata = original ? projectAssistantConnectedAccountTool(original) : undefined - return ( - !metadata?.personalToken && - metadata?.oauth?.required && - providerIdsForService(metadata.oauth.provider).some((providerId) => - providers.has(providerId) - ) + return schemas.filter((schema) => { + const original = getToolMetadata(schema.name) + const metadata = original ? projectAssistantConnectedAccountTool(original) : undefined + return ( + !metadata?.personalToken && + metadata?.oauth?.required && + providerIdsForService(metadata.oauth.provider).some((providerId) => + providers.has(providerId) ) - }) - ) + ) + }) } - return structuredClone(schemas) + return schemas } -async function buildIntegrationToolSchemasUncached({ - userId, - options, - vis, - permissionConfig, -}: IntegrationToolSchemaBuildContext): Promise { - const integrationTools: ToolSchema[] = [] +async function buildIntegrationToolSchemaVariant( + options: IntegrationToolSchemaVariantOptions +): Promise { + const schemas = new Map() + const withEmailTagline = new Map() const { createUserToolSchema } = await import('@/tools/params') - const subscription = await getHighestPrioritySubscription(userId) - const shouldAppendEmailTagline = !subscription || !isPaid(subscription.plan) - const { tools: exposedTools } = projectIntegrationToolsForViewer(vis, permissionConfig) - for (const { toolId, config: toolConfig, service, operation } of exposedTools) { + for (const { toolId, config: toolConfig, service, operation } of getExposedIntegrationTools()) { + await yieldWhenSliceExpires() const metadata = getToolMetadata(toolId) if (options.personalAccountsOnly && !isAssistantIntegrationTool(metadata)) continue const projectedTool = options.personalAccountsOnly @@ -282,37 +295,47 @@ async function buildIntegrationToolSchemasUncached({ } } const catalogEntry = getToolEntry(toolId) - integrationTools.push({ - name: toolId, - service, - operation, - description: getCopilotToolDescription(toolConfig, { + const describe = (appendEmailTagline: boolean) => + getCopilotToolDescription(toolConfig, { isHosted, hostedApiKey: deriveHostedApiKeySupport(toolConfig.hosting), fallbackName: toolId, - appendEmailTagline: shouldAppendEmailTagline, - }), - input_schema: { ...userSchema }, - ...(toolConfig.outputs && { - outputs: Object.fromEntries( - Object.entries(toolConfig.outputs) - .filter(([, output]) => output != null) - .map(([key, output]) => [key, { type: output.type, description: output.description }]) - ), - }), - defer_loading: true, - executeLocally: catalogEntry?.clientExecutable === true || catalogEntry?.route === 'client', - ...(projectedTool.oauth?.required && - isOAuthServiceDeploymentAvailable(projectedTool.oauth.provider) && { - oauth: { - required: true, - provider: projectedTool.oauth.provider, - }, + appendEmailTagline, + }) + const description = describe(false) + // Cloned so freezing the shared schema never freezes registry objects it was built from. + const schema = deepFreeze( + structuredClone({ + name: toolId, + service, + operation, + description, + input_schema: { ...userSchema }, + ...(toolConfig.outputs && { + outputs: Object.fromEntries( + Object.entries(toolConfig.outputs) + .filter(([, output]) => output != null) + .map(([key, output]) => [key, { type: output.type, description: output.description }]) + ), }), - }) + defer_loading: true, + executeLocally: catalogEntry?.clientExecutable === true || catalogEntry?.route === 'client', + ...(projectedTool.oauth?.required && + isOAuthServiceDeploymentAvailable(projectedTool.oauth.provider) && { + oauth: { + required: true, + provider: projectedTool.oauth.provider, + }, + }), + }) + ) + schemas.set(toolId, schema) + const taggedDescription = describe(true) + if (taggedDescription !== description) + withEmailTagline.set(toolId, Object.freeze({ ...schema, description: taggedDescription })) } - return integrationTools + return { schemas, withEmailTagline } } /** diff --git a/apps/sim/lib/mothership/integrations/application/catalog.ts b/apps/sim/lib/mothership/integrations/application/catalog.ts index 8924d8be430..c3b254edba6 100644 --- a/apps/sim/lib/mothership/integrations/application/catalog.ts +++ b/apps/sim/lib/mothership/integrations/application/catalog.ts @@ -107,15 +107,39 @@ function operationTerms(tool: ToolSchema): string[] { return tokenize(`${tool.name} ${tool.service ?? ''} ${tool.description} ${parameters}`) } +interface OperationTermStats { + length: number + frequencies: ReadonlyMap +} + +/** + * Integration schemas are shared, frozen objects reused across requests, so their terms + * are counted once per process; per-request MCP schemas fall out with their request. + */ +const operationTermStats = new WeakMap() + +function getOperationTermStats(tool: ToolSchema): OperationTermStats { + const cached = operationTermStats.get(tool) + if (cached) return cached + const terms = operationTerms(tool) + const frequencies = new Map() + for (const term of terms) frequencies.set(term, (frequencies.get(term) ?? 0) + 1) + const stats = { length: terms.length, frequencies } + operationTermStats.set(tool, stats) + return stats +} + /** Okapi BM25 over the candidate operations; zero means no query term occurs. */ function rankByBm25(tools: ToolSchema[], queryTerms: string[]): number[] { const query = new Set(queryTerms) const documents = tools.map((tool) => { - const terms = operationTerms(tool) + const stats = getOperationTermStats(tool) const frequencies = new Map() - for (const term of terms) - if (query.has(term)) frequencies.set(term, (frequencies.get(term) ?? 0) + 1) - return { length: terms.length, frequencies } + for (const term of query) { + const frequency = stats.frequencies.get(term) + if (frequency) frequencies.set(term, frequency) + } + return { length: stats.length, frequencies } }) const averageLength = documents.reduce((total, document) => total + document.length, 0) / (documents.length || 1) diff --git a/apps/sim/lib/mothership/jobs/budget.ts b/apps/sim/lib/mothership/jobs/budget.ts index 9a38f21f0f1..6cc7b7165a2 100644 --- a/apps/sim/lib/mothership/jobs/budget.ts +++ b/apps/sim/lib/mothership/jobs/budget.ts @@ -1,4 +1,4 @@ -import { db, jobExecutionLogs, workflowSchedule } from '@sim/db' +import { dbFor, jobExecutionLogs, workflowSchedule } from '@sim/db' import { and, eq, sql } from 'drizzle-orm' import type { UsageUpgradePayload } from '@/lib/api/contracts/subscription' @@ -22,7 +22,7 @@ export async function recordJobRunSpend( billingRequestId: string, cumulativeCostUsd: number ): Promise { - const [run] = await db + const [run] = await dbFor('billing') .update(jobExecutionLogs) .set({ cost: sql`jsonb_build_object('total', GREATEST(COALESCE((${jobExecutionLogs.cost}->>'total')::numeric, 0), ${cumulativeCostUsd}::numeric))`, diff --git a/apps/sim/lib/workflows/deployment-status.test.ts b/apps/sim/lib/workflows/deployment-status.test.ts new file mode 100644 index 00000000000..a58fa906328 --- /dev/null +++ b/apps/sim/lib/workflows/deployment-status.test.ts @@ -0,0 +1,77 @@ +/** + * @vitest-environment node + */ +import { workflowDeploymentVersion } from '@sim/db' +import { createMinimalWorkflowState } from '@sim/testing/factories/workflow.factory' +import { + dbChainMock, + dbChainMockFns, + queueTableRows, + resetDbChainMock, +} from '@sim/testing/mocks/database.mock' +import { + workflowsPersistenceUtilsMock, + workflowsPersistenceUtilsMockFns, +} from '@sim/testing/mocks/workflows-persistence-utils.mock' +import { beforeEach, describe, expect, it, vi } from 'vitest' + +vi.mock('@/lib/workflows/persistence/utils', () => workflowsPersistenceUtilsMock) + +import { checkNeedsRedeployment } from '@/lib/workflows/deployment-status' + +const WORKFLOW_ID = 'wf-1' + +describe('checkNeedsRedeployment', () => { + let transactionOpen: boolean + let draftReadWhileTransactionOpen: boolean + + beforeEach(() => { + resetDbChainMock() + transactionOpen = false + draftReadWhileTransactionOpen = false + + dbChainMockFns.transaction.mockImplementation(async (cb: (tx: unknown) => unknown) => { + transactionOpen = true + try { + return await cb(dbChainMock.db) + } finally { + transactionOpen = false + } + }) + + queueTableRows(workflowDeploymentVersion, [ + { id: 'version-1', state: { blocks: {} }, workspaceId: 'ws-1' }, + ]) + + const deployed = createMinimalWorkflowState() + const draft = createMinimalWorkflowState() + draft.blocks.agent1.subBlocks.prompt.value = 'Edited after deploy' + const draftBlocks = draft.blocks + + workflowsPersistenceUtilsMockFns.mockMaterializeDeploymentState.mockResolvedValue({ + ...deployed, + parallels: {}, + variables: {}, + }) + workflowsPersistenceUtilsMockFns.mockLoadWorkflowDeploymentSnapshot.mockResolvedValue({ + edges: draft.edges, + loops: {}, + parallels: {}, + variables: {}, + get blocks() { + if (transactionOpen) draftReadWhileTransactionOpen = true + return draftBlocks + }, + }) + }) + + it('diffs the draft against the deployment only after the snapshot transaction commits', async () => { + await expect(checkNeedsRedeployment(WORKFLOW_ID)).resolves.toBe(true) + expect(draftReadWhileTransactionOpen).toBe(false) + }) + + it('reports no redeploy when the workflow has no active deployment', async () => { + resetDbChainMock() + await expect(checkNeedsRedeployment(WORKFLOW_ID)).resolves.toBe(false) + }) +}) diff --git a/apps/sim/lib/workflows/deployment-status.ts b/apps/sim/lib/workflows/deployment-status.ts index 5931625a7d1..df2b15e8d58 100644 --- a/apps/sim/lib/workflows/deployment-status.ts +++ b/apps/sim/lib/workflows/deployment-status.ts @@ -19,9 +19,12 @@ import type { WorkflowState } from '@/stores/workflows/workflow/types' * caller is what let this surface compare a raw jsonb blob against a normalized * draft, so that the server and the client answered the same question * differently for the same workflow. + * + * All reads share one REPEATABLE READ snapshot; the CPU-only diff runs after + * commit so the connection is not held idle in transaction. */ export async function checkNeedsRedeployment(workflowId: string): Promise { - return db.transaction(async (tx) => { + const snapshot = await db.transaction(async (tx) => { await tx.execute(sql`SET TRANSACTION ISOLATION LEVEL REPEATABLE READ`) /* * `workspaceId` is selected here, in this transaction, rather than left for @@ -53,10 +56,10 @@ export async function checkNeedsRedeployment(workflowId: string): Promise { if (!workspaceId) { return null } - const rows = await db + const rows = await executor .select({ billedAccountUserId: workspaceTable.billedAccountUserId, allowPersonalApiKeys: workspaceTable.allowPersonalApiKeys, diff --git a/packages/db/db.test.ts b/packages/db/db.test.ts new file mode 100644 index 00000000000..851e5baedb5 --- /dev/null +++ b/packages/db/db.test.ts @@ -0,0 +1,36 @@ +import { beforeAll, describe, expect, it, vi } from 'vitest' + +vi.mock('@sim/logger', () => ({ + createLogger: () => ({ error: vi.fn(), warn: vi.fn(), info: vi.fn(), debug: vi.fn() }), +})) + +type DbModule = typeof import('./db') + +let database: DbModule + +/** postgres.js connects on first query, so unreachable URLs build real clients without a server. */ +beforeAll(async () => { + vi.stubEnv('SIM_DB_ROLE', undefined) + vi.stubEnv('DATABASE_URL_WEB', undefined) + vi.stubEnv('DATABASE_REPLICA_URL_WEB', undefined) + vi.stubEnv('DATABASE_URL', 'postgres://127.0.0.1:1/primary') + vi.stubEnv('DATABASE_REPLICA_URL', 'postgres://127.0.0.1:2/replica') + database = await import('./db') +}) + +describe('isPrimaryPoolClient', () => { + it('recognizes sub-pool clients, so process-wide caches keep serving their reads', () => { + expect(database.isPrimaryPoolClient(database.db)).toBe(true) + expect(database.isPrimaryPoolClient(database.dbFor('auth'))).toBe(true) + expect(database.isPrimaryPoolClient(database.dbFor('billing'))).toBe(true) + }) + + it('rejects a separately configured replica, whose reads can lag the primary', () => { + expect(database.dbReplica).not.toBe(database.db) + expect(database.isPrimaryPoolClient(database.dbReplica)).toBe(false) + }) + + it('rejects an executor that is not one of the root clients', () => { + expect(database.isPrimaryPoolClient({ ...database.db })).toBe(false) + }) +}) diff --git a/packages/db/db.ts b/packages/db/db.ts index 7176f92c6f1..c4a4f09f1ff 100644 --- a/packages/db/db.ts +++ b/packages/db/db.ts @@ -74,6 +74,19 @@ const postgresClient = instrumentPoolClient( export const db = drizzle(postgresClient, { schema }) +/** Root clients connected to the primary: `db` and every `dbFor()` sub-pool. */ +const primaryPoolClients = new WeakSet([db]) + +/** + * True for a root client on the primary — `db` or any `dbFor()` sub-pool — and + * false for a transaction handle or `dbReplica`. A process-wide cache may only + * serve and store reads made through such a client: a transaction can see its + * own uncommitted writes, and the replica can lag. + */ +export function isPrimaryPoolClient(executor: object): boolean { + return primaryPoolClients.has(executor) +} + /** * Opt-in read-replica client for reads that tolerate bounded staleness and have * no read-your-writes dependency (logs, exports, dashboard aggregations). Never @@ -159,5 +172,6 @@ export function dbFor(role: SubProcessDbRole): typeof db { { schema } ) subPoolClients.set(role, client) + primaryPoolClients.add(client) return client } diff --git a/packages/db/pool-profiles.ts b/packages/db/pool-profiles.ts index 7ffb19bd73a..4ce1a66c9ba 100644 --- a/packages/db/pool-profiles.ts +++ b/packages/db/pool-profiles.ts @@ -8,4 +8,12 @@ export const DB_POOL_PROFILES = { cleanup: { primaryMax: 5, replicaMax: 2, appName: 'sim-cleanup' }, exec: { primaryMax: 10, replicaMax: 4, appName: 'sim-exec' }, search: { primaryMax: 5, replicaMax: 0, appName: 'sim-search' }, + /** + * Usage-ledger writes and threshold settlements: short transactions bounded by + * server-side timeouts, so they must never queue behind slow reads on the process + * pool, where postgres.js pipelines new work onto busy connections. + */ + billing: { primaryMax: 4, replicaMax: 0, appName: 'sim-billing' }, + /** Session, adapter, and API-key lookups that gate every authenticated request. */ + auth: { primaryMax: 3, replicaMax: 0, appName: 'sim-auth' }, } as const diff --git a/packages/testing/src/mocks/database.mock.ts b/packages/testing/src/mocks/database.mock.ts index 8198651e206..f6eb2535d83 100644 --- a/packages/testing/src/mocks/database.mock.ts +++ b/packages/testing/src/mocks/database.mock.ts @@ -439,6 +439,8 @@ const dbInstance = { transaction, } +const mockIsPrimaryPoolClient = vi.fn((executor: object): boolean => executor === dbInstance) + /** * Static mock module for `@sim/db` backed by `dbChainMockFns`. * @@ -451,8 +453,9 @@ export const dbChainMock = { db: dbInstance, /** Same instance as `db` so per-test chain overrides cover both clients. */ dbReplica: dbInstance, - /** Sub-pool clients (`dbFor('cleanup' | 'exec')`) share the same instance too. */ + /** Sub-pool clients (`dbFor('auth' | 'billing' | ...)`) share the same instance too. */ dbFor: () => dbInstance, + isPrimaryPoolClient: mockIsPrimaryPoolClient, runOutsideTransactionContext: (fn: () => T): T => fn(), instrumentPoolClient: (client: T): T => client, } @@ -474,6 +477,8 @@ function resolveDbUrl( */ export const databaseMockFns = { mockResolveDbUrl: vi.fn(resolveDbUrl), + /** Recognizes only the shared instance unless a test marks its own client. */ + mockIsPrimaryPoolClient, mockEnsureRowCountTriggers: vi.fn(async (): Promise => {}), mockVerifyRowCountTriggers: vi.fn(async () => ({ incrementTrigger: true, diff --git a/packages/testing/src/mocks/schema-tables.generated.ts b/packages/testing/src/mocks/schema-tables.generated.ts index de7a334e001..afa1de7bab1 100644 --- a/packages/testing/src/mocks/schema-tables.generated.ts +++ b/packages/testing/src/mocks/schema-tables.generated.ts @@ -2339,6 +2339,16 @@ export const GENERATED_DB_CONSTANTS = { replicaMax: 0, appName: 'sim-search', }, + billing: { + primaryMax: 4, + replicaMax: 0, + appName: 'sim-billing', + }, + auth: { + primaryMax: 3, + replicaMax: 0, + appName: 'sim-auth', + }, }, INCREMENT_ROW_COUNT_SQL: '\nCREATE OR REPLACE FUNCTION increment_table_row_count()\nRETURNS TRIGGER AS $$\nBEGIN\n UPDATE user_table_definitions\n SET row_count = row_count + 1\n WHERE id = NEW.table_id;\n RETURN NEW;\nEND;\n$$ LANGUAGE plpgsql;\n\nDROP TRIGGER IF EXISTS trg_increment_row_count ON user_table_rows;\nCREATE TRIGGER trg_increment_row_count\nAFTER INSERT ON user_table_rows\nFOR EACH ROW\nEXECUTE FUNCTION increment_table_row_count();\n',