From 935d0dc9b653614b38ca6cbbfcddc75c3e11504f Mon Sep 17 00:00:00 2001 From: Roomote Date: Tue, 25 Aug 2026 13:46:50 +0000 Subject: [PATCH 1/4] refactor: centralize integration lifecycle transitions --- .../callback/__tests__/route.test.ts | 28 +- .../src/app/api/mcp-oauth/callback/route.ts | 64 +- .../lib/server/deployment-mcp-connection.ts | 202 ++++ .../src/lib/server/integration-telemetry.ts | 22 + apps/web/src/lib/server/mcp-linear.test.ts | 1 - apps/web/src/lib/server/mcp-linear.ts | 26 +- .../trpc/commands/custom-mcp-servers/index.ts | 21 +- .../commands/mcp-connections/index.test.ts | 31 + .../trpc/commands/mcp-connections/index.ts | 1055 ++++------------- 9 files changed, 562 insertions(+), 888 deletions(-) create mode 100644 apps/web/src/lib/server/deployment-mcp-connection.ts diff --git a/apps/web/src/app/api/mcp-oauth/callback/__tests__/route.test.ts b/apps/web/src/app/api/mcp-oauth/callback/__tests__/route.test.ts index 2a94b96dd..d449ae9b3 100644 --- a/apps/web/src/app/api/mcp-oauth/callback/__tests__/route.test.ts +++ b/apps/web/src/app/api/mcp-oauth/callback/__tests__/route.test.ts @@ -307,7 +307,6 @@ describe('GET /api/mcp-oauth/callback', () => { refresh_token: 'refresh-token', }, replayToken: null, - enabledByUserId: 'user-1', }); expect(storeTokensMock).not.toHaveBeenCalled(); }); @@ -433,6 +432,33 @@ describe('GET /api/mcp-oauth/callback', () => { }); }); + it('captures enablement when a concurrent disable wins the insert race', async () => { + mcpConnectionsFindFirstMock.mockResolvedValue({ + id: CONNECTION_ID, + mcpId: 'resend', + userId: null, + connectionRole: 'default', + }); + getMcpIntegrationMock.mockReturnValue({ + id: 'resend', + name: 'Resend', + url: 'https://mcp.resend.com/mcp', + }); + isDeploymentScopedMcpIntegrationMock.mockReturnValue(true); + deploymentEnablementUpdateReturningMock + .mockResolvedValueOnce([]) + .mockResolvedValueOnce([{ mcpId: 'resend' }]); + deploymentEnablementInsertReturningMock.mockResolvedValue([]); + + await GET(buildRequest('?code=auth-code&state=state-1')); + + expect(deploymentEnablementUpdateReturningMock).toHaveBeenCalledTimes(2); + expect(captureEventMock).toHaveBeenCalledWith('integration_enabled', { + userId: 'user-1', + properties: { integration_id: 'resend' }, + }); + }); + it('surfaces token exchange failures with a safe reason and stage', async () => { exchangeCodeForTokensMock.mockRejectedValueOnce( new Error('provider response omitted'), diff --git a/apps/web/src/app/api/mcp-oauth/callback/route.ts b/apps/web/src/app/api/mcp-oauth/callback/route.ts index f40131eb3..849f402e6 100644 --- a/apps/web/src/app/api/mcp-oauth/callback/route.ts +++ b/apps/web/src/app/api/mcp-oauth/callback/route.ts @@ -1,13 +1,7 @@ import { createHash } from 'node:crypto'; import { type NextRequest, NextResponse } from 'next/server'; -import { - and, - db, - mcpConnections, - deploymentMcpEnablements, - eq, -} from '@roomote/db/server'; +import { db, mcpConnections, eq } from '@roomote/db/server'; import { getMcpIntegration, getMcpIntegrationDefaultDisabledTools, @@ -30,7 +24,8 @@ import { authorize } from '@/lib/server'; import { bootstrapWebRuntimeEnv } from '@/lib/server/bootstrap-runtime-env'; import { getPublicAppUrl } from '@/lib/server/get-public-app-url'; import { logger } from '@/lib/server/logger'; -import { captureIntegrationLifecycleEvent } from '@/lib/server/integration-telemetry'; +import { captureIntegrationConnectionTransitions } from '@/lib/server/integration-telemetry'; +import { enableDeploymentMcpIntegration } from '@/lib/server/deployment-mcp-connection'; import { hydrateLinearMcpConnectionAfterOauth, LinearReplayIdentityMismatchError, @@ -367,7 +362,6 @@ export async function GET(request: NextRequest) { connection, tokens, replayToken: oauthState.replayToken, - enabledByUserId: userId, }); } else { failureStage = 'token_storage'; @@ -381,51 +375,19 @@ export async function GET(request: NextRequest) { failureStage = 'deployment_enablement'; const defaultDisabledTools = getMcpIntegrationDefaultDisabledTools(integration); - const [reenabled] = await db - .update(deploymentMcpEnablements) - .set({ - enabled: true, - enabledByUserId: userId, - updatedAt: new Date(), - }) - .where( - and( - eq(deploymentMcpEnablements.mcpId, integration.id), - eq(deploymentMcpEnablements.enabled, false), - ), - ) - .returning({ mcpId: deploymentMcpEnablements.mcpId }); - const [inserted] = reenabled - ? [] - : await db - .insert(deploymentMcpEnablements) - .values({ - mcpId: integration.id, - enabled: true, - enabledByUserId: userId, - ...(defaultDisabledTools.length > 0 - ? { - disabledTools: [...defaultDisabledTools], - } - : {}), - }) - .onConflictDoNothing({ target: deploymentMcpEnablements.mcpId }) - .returning({ mcpId: deploymentMcpEnablements.mcpId }); - integrationBecameEnabled = Boolean(reenabled || inserted); + integrationBecameEnabled = await enableDeploymentMcpIntegration({ + mcpId: integration.id, + userId, + defaultDisabledTools: [...defaultDisabledTools], + }); } - captureIntegrationLifecycleEvent( - 'integration_connected', - connection.mcpId, + captureIntegrationConnectionTransitions({ + integrationId: connection.mcpId, userId, - ); - if (integrationBecameEnabled && integration) { - captureIntegrationLifecycleEvent( - 'integration_enabled', - integration.id, - userId, - ); - } + connected: true, + enabled: integrationBecameEnabled, + }); return redirectToResult({ status: 'connected' }); } catch (error) { diff --git a/apps/web/src/lib/server/deployment-mcp-connection.ts b/apps/web/src/lib/server/deployment-mcp-connection.ts new file mode 100644 index 000000000..aac88fc41 --- /dev/null +++ b/apps/web/src/lib/server/deployment-mcp-connection.ts @@ -0,0 +1,202 @@ +import { + and, + db, + deploymentMcpEnablements, + eq, + isNull, + mcpConnections, + type DatabaseOrTransaction, +} from '@roomote/db/server'; +import type { McpConnectionAuthConfig } from '@roomote/types'; + +import { captureIntegrationConnectionTransitions } from '@/lib/server/integration-telemetry'; + +type SaveAuthenticatedDeploymentMcpConnectionInput< + TAuthConfig extends McpConnectionAuthConfig, +> = { + mcpId: string; + userId: string; + buildAuthConfig: ( + existingAuthConfig: McpConnectionAuthConfig | null, + ) => TAuthConfig | Promise; + clearOauthTokens?: boolean; + resetDisabledTools?: boolean; +}; + +async function persistDeploymentMcpIntegrationEnabled( + database: DatabaseOrTransaction, + input: { + mcpId: string; + userId: string; + defaultDisabledTools?: string[]; + resetDisabledTools?: boolean; + }, +): Promise { + const updatedDisabledTools = input.resetDisabledTools ? null : undefined; + const insertedDisabledTools = input.resetDisabledTools + ? null + : input.defaultDisabledTools?.length + ? input.defaultDisabledTools + : undefined; + const [reenabled] = await database + .update(deploymentMcpEnablements) + .set({ + enabled: true, + enabledByUserId: input.userId, + ...(updatedDisabledTools !== undefined + ? { disabledTools: updatedDisabledTools } + : {}), + updatedAt: new Date(), + }) + .where( + and( + eq(deploymentMcpEnablements.mcpId, input.mcpId), + eq(deploymentMcpEnablements.enabled, false), + ), + ) + .returning({ mcpId: deploymentMcpEnablements.mcpId }); + + if (reenabled) { + return true; + } + + const [inserted] = await database + .insert(deploymentMcpEnablements) + .values({ + mcpId: input.mcpId, + enabled: true, + enabledByUserId: input.userId, + ...(insertedDisabledTools !== undefined + ? { disabledTools: insertedDisabledTools } + : {}), + }) + .onConflictDoNothing({ target: deploymentMcpEnablements.mcpId }) + .returning({ mcpId: deploymentMcpEnablements.mcpId }); + + if (inserted) { + return true; + } + + const [enabledAfterConflict] = await database + .update(deploymentMcpEnablements) + .set({ + enabled: true, + enabledByUserId: input.userId, + ...(updatedDisabledTools !== undefined + ? { disabledTools: updatedDisabledTools } + : {}), + updatedAt: new Date(), + }) + .where( + and( + eq(deploymentMcpEnablements.mcpId, input.mcpId), + eq(deploymentMcpEnablements.enabled, false), + ), + ) + .returning({ mcpId: deploymentMcpEnablements.mcpId }); + + if (enabledAfterConflict) { + return true; + } + + await database + .update(deploymentMcpEnablements) + .set({ + enabledByUserId: input.userId, + ...(updatedDisabledTools !== undefined + ? { disabledTools: updatedDisabledTools } + : {}), + updatedAt: new Date(), + }) + .where(eq(deploymentMcpEnablements.mcpId, input.mcpId)); + + return false; +} + +export async function saveAuthenticatedDeploymentMcpConnection< + TAuthConfig extends McpConnectionAuthConfig, +>( + input: SaveAuthenticatedDeploymentMcpConnectionInput, +): Promise { + const existingConnection = await db.query.mcpConnections.findFirst({ + where: and( + eq(mcpConnections.mcpId, input.mcpId), + isNull(mcpConnections.userId), + ), + columns: { authConfig: true }, + }); + const authConfig = await input.buildAuthConfig( + existingConnection?.authConfig ?? null, + ); + + const transitions = await db.transaction(async (tx) => { + const [insertedConnection] = await tx + .insert(mcpConnections) + .values({ + userId: null, + mcpId: input.mcpId, + connectionRole: 'default', + authConfig, + enabled: true, + authStatus: 'authenticated', + }) + .onConflictDoNothing({ + target: [ + mcpConnections.userId, + mcpConnections.mcpId, + mcpConnections.connectionRole, + ], + }) + .returning({ id: mcpConnections.id }); + + if (!insertedConnection) { + await tx + .update(mcpConnections) + .set({ + authConfig, + ...(input.clearOauthTokens + ? { + accessToken: null, + refreshToken: null, + tokenExpiresAt: null, + scopes: null, + } + : {}), + enabled: true, + authStatus: 'authenticated', + updatedAt: new Date(), + }) + .where( + and( + eq(mcpConnections.mcpId, input.mcpId), + isNull(mcpConnections.userId), + eq(mcpConnections.connectionRole, 'default'), + ), + ); + } + + const enabled = await persistDeploymentMcpIntegrationEnabled(tx, { + mcpId: input.mcpId, + userId: input.userId, + resetDisabledTools: input.resetDisabledTools, + }); + + return { connected: Boolean(insertedConnection), enabled }; + }); + + captureIntegrationConnectionTransitions({ + integrationId: input.mcpId, + userId: input.userId, + ...transitions, + }); + + return authConfig; +} + +export async function enableDeploymentMcpIntegration(input: { + mcpId: string; + userId: string; + defaultDisabledTools: string[]; +}): Promise { + return persistDeploymentMcpIntegrationEnabled(db, input); +} diff --git a/apps/web/src/lib/server/integration-telemetry.ts b/apps/web/src/lib/server/integration-telemetry.ts index b24efe7ec..8b37ae8d9 100644 --- a/apps/web/src/lib/server/integration-telemetry.ts +++ b/apps/web/src/lib/server/integration-telemetry.ts @@ -16,3 +16,25 @@ export function captureIntegrationLifecycleEvent( properties: { integration_id: integrationId }, }); } + +export function captureIntegrationConnectionTransitions(input: { + integrationId: string; + userId: string; + connected: boolean; + enabled: boolean; +}): void { + if (input.connected) { + captureIntegrationLifecycleEvent( + 'integration_connected', + input.integrationId, + input.userId, + ); + } + if (input.enabled) { + captureIntegrationLifecycleEvent( + 'integration_enabled', + input.integrationId, + input.userId, + ); + } +} diff --git a/apps/web/src/lib/server/mcp-linear.test.ts b/apps/web/src/lib/server/mcp-linear.test.ts index fe331efc5..3ee1b7681 100644 --- a/apps/web/src/lib/server/mcp-linear.test.ts +++ b/apps/web/src/lib/server/mcp-linear.test.ts @@ -61,7 +61,6 @@ vi.mock('@roomote/db/server', () => ({ query: { mcpConnections: { findFirst: vi.fn() } }, update: dbUpdateMock, }, - deploymentMcpEnablements: {}, eq: vi.fn(), mcpConnections: { id: 'id' }, })); diff --git a/apps/web/src/lib/server/mcp-linear.ts b/apps/web/src/lib/server/mcp-linear.ts index 208c1b0e9..50323c287 100644 --- a/apps/web/src/lib/server/mcp-linear.ts +++ b/apps/web/src/lib/server/mcp-linear.ts @@ -1,12 +1,7 @@ import { LinearClient } from '@linear/sdk'; import { buildTaskStartingText } from '@roomote/communication/chat-messages'; -import { - db, - mcpConnections, - eq, - deploymentMcpEnablements, -} from '@roomote/db/server'; +import { db, mcpConnections, eq } from '@roomote/db/server'; import { consumeMcpOauthReplay, findLinearDeploymentMcpConnection, @@ -214,7 +209,6 @@ export async function hydrateLinearMcpConnectionAfterOauth(input: { connection: NonNullable; tokens: OAuthTokens; replayToken?: string | null; - enabledByUserId?: string; }) { const viewerClient = new LinearClient({ accessToken: input.tokens.access_token, @@ -236,24 +230,6 @@ export async function hydrateLinearMcpConnectionAfterOauth(input: { appUserId: viewer.id, }); - if (input.enabledByUserId) { - await db - .insert(deploymentMcpEnablements) - .values({ - mcpId: 'linear', - enabled: true, - enabledByUserId: input.enabledByUserId, - }) - .onConflictDoUpdate({ - target: deploymentMcpEnablements.mcpId, - set: { - enabled: true, - enabledByUserId: input.enabledByUserId, - updatedAt: new Date(), - }, - }); - } - return; } diff --git a/apps/web/src/trpc/commands/custom-mcp-servers/index.ts b/apps/web/src/trpc/commands/custom-mcp-servers/index.ts index 9731217f3..09f8286c0 100644 --- a/apps/web/src/trpc/commands/custom-mcp-servers/index.ts +++ b/apps/web/src/trpc/commands/custom-mcp-servers/index.ts @@ -20,7 +20,10 @@ import { } from '@roomote/types'; import type { UserAuthSuccess } from '@/types'; import { Env, isCustomMcpDisabled } from '@/lib/server/env'; -import { captureIntegrationLifecycleEvent } from '@/lib/server/integration-telemetry'; +import { + captureIntegrationConnectionTransitions, + captureIntegrationLifecycleEvent, +} from '@/lib/server/integration-telemetry'; export const CUSTOM_MCP_DISABLED_MESSAGE = 'Custom MCP servers are disabled by the deployment operator.'; @@ -242,18 +245,12 @@ export async function createCustomMcpServerCommand( ); const integrationId = customMcpConnectionId(created.id); - captureIntegrationLifecycleEvent( - 'integration_enabled', + captureIntegrationConnectionTransitions({ integrationId, - auth.userId, - ); - if (input.transport === 'stdio' || input.authType !== 'oauth') { - captureIntegrationLifecycleEvent( - 'integration_connected', - integrationId, - auth.userId, - ); - } + userId: auth.userId, + connected: input.transport === 'stdio' || input.authType !== 'oauth', + enabled: true, + }); return { id: created.id }; } diff --git a/apps/web/src/trpc/commands/mcp-connections/index.test.ts b/apps/web/src/trpc/commands/mcp-connections/index.test.ts index 797f7bcc3..fb7ba02ab 100644 --- a/apps/web/src/trpc/commands/mcp-connections/index.test.ts +++ b/apps/web/src/trpc/commands/mcp-connections/index.test.ts @@ -1,6 +1,7 @@ import { db, deploymentMcpEnablements, + eq, mcpConnections, userFactory, } from '@roomote/db/server'; @@ -85,4 +86,34 @@ describe('MCP connection lifecycle telemetry', () => { expect(captureEventMock).not.toHaveBeenCalled(); }); + + it('captures credential connection and enablement transitions independently', async () => { + await db.insert(deploymentMcpEnablements).values({ + mcpId: 'asana', + enabled: true, + enabledByUserId: adminAuth.userId, + }); + + await saveAsanaConnectionCommand(adminAuth, { accessToken: 'asana-token' }); + + expect(captureEventMock).toHaveBeenCalledTimes(1); + expect(captureEventMock).toHaveBeenCalledWith('integration_connected', { + userId: adminAuth.userId, + properties: { integration_id: 'asana' }, + }); + + captureEventMock.mockClear(); + await db + .update(deploymentMcpEnablements) + .set({ enabled: false }) + .where(eq(deploymentMcpEnablements.mcpId, 'asana')); + + await saveAsanaConnectionCommand(adminAuth, { accessToken: '' }); + + expect(captureEventMock).toHaveBeenCalledTimes(1); + expect(captureEventMock).toHaveBeenCalledWith('integration_enabled', { + userId: adminAuth.userId, + properties: { integration_id: 'asana' }, + }); + }); }); diff --git a/apps/web/src/trpc/commands/mcp-connections/index.ts b/apps/web/src/trpc/commands/mcp-connections/index.ts index 0283e5731..f2f95dece 100644 --- a/apps/web/src/trpc/commands/mcp-connections/index.ts +++ b/apps/web/src/trpc/commands/mcp-connections/index.ts @@ -45,6 +45,7 @@ import { Env, areCuratedIntegrationsDisabled } from '@/lib/server/env'; import { assertCuratedIntegrationsEnabled } from '@/lib/server/curated-integrations'; import { MCP_TOOL_CATALOG_REQUIRES_PERSONAL_CONNECTION } from '@/lib/mcp-tool-errors'; import { captureIntegrationLifecycleEvent } from '@/lib/server/integration-telemetry'; +import { saveAuthenticatedDeploymentMcpConnection } from '@/lib/server/deployment-mcp-connection'; import type { SaveAsanaConnectionInput, SaveNotionConnectionInput, @@ -982,152 +983,97 @@ export async function saveSnowflakeConnectionCommand( assertAdmin(auth); assertCuratedIntegrationsEnabled(); - const existingConnection = await db.query.mcpConnections.findFirst({ - where: and( - eq(mcpConnections.mcpId, 'snowflake'), - isNull(mcpConnections.userId), - ), - columns: { - authConfig: true, - }, - }); - - const existingConfig = isMcpConnectionSnowflakeConfig( - existingConnection?.authConfig, - ) - ? existingConnection.authConfig - : null; - const preservingExistingCredential = - input.authMethod === 'password' && input.password.length === 0; - const nextEncryptedPassword = - input.authMethod === 'password' - ? input.password.length > 0 - ? encrypt(input.password) - : existingConfig?.encryptedPassword - : undefined; - const nextEncryptedPrivateKey = - input.authMethod === 'key_pair' - ? input.privateKey.trim().length > 0 - ? encrypt(input.privateKey) - : existingConfig?.encryptedPrivateKey - : preservingExistingCredential - ? existingConfig?.encryptedPrivateKey - : undefined; - const nextEncryptedPrivateKeyPassphrase = - input.authMethod === 'key_pair' - ? input.privateKeyPassphrase.length > 0 - ? encrypt(input.privateKeyPassphrase) - : input.privateKey.trim().length > 0 - ? undefined - : existingConfig?.encryptedPrivateKeyPassphrase - : preservingExistingCredential - ? existingConfig?.encryptedPrivateKeyPassphrase - : undefined; - - const authConfig = { - type: 'snowflake' as const, - account: input.account, - username: input.username, - role: input.role, - ...(input.warehouse - ? { warehouse: input.warehouse } - : existingConfig?.warehouse - ? { warehouse: existingConfig.warehouse } - : {}), - ...(input.database - ? { database: input.database } - : existingConfig?.database - ? { database: existingConfig.database } - : {}), - ...(nextEncryptedPassword - ? { encryptedPassword: nextEncryptedPassword } - : {}), - ...(nextEncryptedPrivateKey - ? { encryptedPrivateKey: nextEncryptedPrivateKey } - : {}), - ...(nextEncryptedPrivateKeyPassphrase - ? { - encryptedPrivateKeyPassphrase: nextEncryptedPrivateKeyPassphrase, - } - : {}), - ...(existingConfig?.schema ? { schema: existingConfig.schema } : {}), - ...(existingConfig?.allowedStatementTypes - ? { allowedStatementTypes: existingConfig.allowedStatementTypes } - : {}), - }; - - if ( - input.authMethod === 'password' && - !authConfig.encryptedPassword && - !authConfig.encryptedPrivateKey - ) { - throw new Error( - 'Programmatic Access Token is required when no Snowflake credential is already stored.', - ); - } - - if (input.authMethod === 'key_pair' && !authConfig.encryptedPrivateKey) { - throw new Error( - 'Private key is required when no Snowflake key pair is already stored.', - ); - } - - await db - .insert(mcpConnections) - .values({ - userId: null, - mcpId: 'snowflake', - connectionRole: 'default', - authConfig, - enabled: true, - authStatus: 'authenticated', - }) - .onConflictDoUpdate({ - target: [ - mcpConnections.userId, - mcpConnections.mcpId, - mcpConnections.connectionRole, - ], - set: { - connectionRole: 'default', - authConfig, - enabled: true, - authStatus: 'authenticated', - updatedAt: new Date(), - }, - }); + const authConfig = await saveAuthenticatedDeploymentMcpConnection({ + mcpId: 'snowflake', + userId: auth.userId, + buildAuthConfig: (existingAuthConfig) => { + const existingConfig = isMcpConnectionSnowflakeConfig(existingAuthConfig) + ? existingAuthConfig + : null; + const preservingExistingCredential = + input.authMethod === 'password' && input.password.length === 0; + const nextEncryptedPassword = + input.authMethod === 'password' + ? input.password.length > 0 + ? encrypt(input.password) + : existingConfig?.encryptedPassword + : undefined; + const nextEncryptedPrivateKey = + input.authMethod === 'key_pair' + ? input.privateKey.trim().length > 0 + ? encrypt(input.privateKey) + : existingConfig?.encryptedPrivateKey + : preservingExistingCredential + ? existingConfig?.encryptedPrivateKey + : undefined; + const nextEncryptedPrivateKeyPassphrase = + input.authMethod === 'key_pair' + ? input.privateKeyPassphrase.length > 0 + ? encrypt(input.privateKeyPassphrase) + : input.privateKey.trim().length > 0 + ? undefined + : existingConfig?.encryptedPrivateKeyPassphrase + : preservingExistingCredential + ? existingConfig?.encryptedPrivateKeyPassphrase + : undefined; + + const nextAuthConfig = { + type: 'snowflake' as const, + account: input.account, + username: input.username, + role: input.role, + ...(input.warehouse + ? { warehouse: input.warehouse } + : existingConfig?.warehouse + ? { warehouse: existingConfig.warehouse } + : {}), + ...(input.database + ? { database: input.database } + : existingConfig?.database + ? { database: existingConfig.database } + : {}), + ...(nextEncryptedPassword + ? { encryptedPassword: nextEncryptedPassword } + : {}), + ...(nextEncryptedPrivateKey + ? { encryptedPrivateKey: nextEncryptedPrivateKey } + : {}), + ...(nextEncryptedPrivateKeyPassphrase + ? { + encryptedPrivateKeyPassphrase: nextEncryptedPrivateKeyPassphrase, + } + : {}), + ...(existingConfig?.schema ? { schema: existingConfig.schema } : {}), + ...(existingConfig?.allowedStatementTypes + ? { allowedStatementTypes: existingConfig.allowedStatementTypes } + : {}), + }; + + if ( + input.authMethod === 'password' && + !nextAuthConfig.encryptedPassword && + !nextAuthConfig.encryptedPrivateKey + ) { + throw new Error( + 'Programmatic Access Token is required when no Snowflake credential is already stored.', + ); + } - await db - .insert(deploymentMcpEnablements) - .values({ - mcpId: 'snowflake', - enabled: true, - enabledByUserId: auth.userId, - }) - .onConflictDoUpdate({ - target: [deploymentMcpEnablements.mcpId], - set: { - enabled: true, - enabledByUserId: auth.userId, - updatedAt: new Date(), - }, - }); + if ( + input.authMethod === 'key_pair' && + !nextAuthConfig.encryptedPrivateKey + ) { + throw new Error( + 'Private key is required when no Snowflake key pair is already stored.', + ); + } - if (!existingConnection) { - captureIntegrationLifecycleEvent( - 'integration_connected', - 'snowflake', - auth.userId, - ); - captureIntegrationLifecycleEvent( - 'integration_enabled', - 'snowflake', - auth.userId, - ); - } + return nextAuthConfig; + }, + }); return { - authMethod: nextEncryptedPrivateKey + authMethod: authConfig.encryptedPrivateKey ? ('key_pair' as const) : ('password' as const), account: authConfig.account, @@ -1145,91 +1091,27 @@ export async function saveAsanaConnectionCommand( assertAdmin(auth); assertCuratedIntegrationsEnabled(); - const existingConnection = await db.query.mcpConnections.findFirst({ - where: and( - eq(mcpConnections.mcpId, 'asana'), - isNull(mcpConnections.userId), - ), - columns: { - authConfig: true, + await saveAuthenticatedDeploymentMcpConnection({ + mcpId: 'asana', + userId: auth.userId, + buildAuthConfig: (existingAuthConfig) => { + const existingConfig = isMcpConnectionAsanaConfig(existingAuthConfig) + ? existingAuthConfig + : null; + const encryptedToken = input.accessToken.length + ? encrypt(input.accessToken) + : existingConfig?.encryptedToken; + + if (!encryptedToken) { + throw new Error( + 'Asana access token is required when no Asana token is already stored.', + ); + } + + return { type: 'asana' as const, encryptedToken }; }, }); - const existingConfig = isMcpConnectionAsanaConfig( - existingConnection?.authConfig, - ) - ? existingConnection.authConfig - : null; - const nextEncryptedToken = - input.accessToken.length > 0 - ? encrypt(input.accessToken) - : existingConfig?.encryptedToken; - - if (!nextEncryptedToken) { - throw new Error( - 'Asana access token is required when no Asana token is already stored.', - ); - } - - const authConfig = { - type: 'asana' as const, - encryptedToken: nextEncryptedToken, - }; - - await db - .insert(mcpConnections) - .values({ - userId: null, - mcpId: 'asana', - connectionRole: 'default', - authConfig, - enabled: true, - authStatus: 'authenticated', - }) - .onConflictDoUpdate({ - target: [ - mcpConnections.userId, - mcpConnections.mcpId, - mcpConnections.connectionRole, - ], - set: { - connectionRole: 'default', - authConfig, - enabled: true, - authStatus: 'authenticated', - updatedAt: new Date(), - }, - }); - - await db - .insert(deploymentMcpEnablements) - .values({ - mcpId: 'asana', - enabled: true, - enabledByUserId: auth.userId, - }) - .onConflictDoUpdate({ - target: [deploymentMcpEnablements.mcpId], - set: { - enabled: true, - enabledByUserId: auth.userId, - updatedAt: new Date(), - }, - }); - - if (!existingConnection) { - captureIntegrationLifecycleEvent( - 'integration_connected', - 'asana', - auth.userId, - ); - captureIntegrationLifecycleEvent( - 'integration_enabled', - 'asana', - auth.userId, - ); - } - return { authStatus: 'authenticated' as const, }; @@ -1242,92 +1124,28 @@ export async function saveNotionConnectionCommand( assertAdmin(auth); assertCuratedIntegrationsEnabled(); - const existingConnection = await db.query.mcpConnections.findFirst({ - where: and( - eq(mcpConnections.mcpId, 'notion'), - isNull(mcpConnections.userId), - ), - columns: { authConfig: true }, - }); - const existingConfig = isMcpConnectionNotionConfig( - existingConnection?.authConfig, - ) - ? existingConnection.authConfig - : null; - const nextEncryptedToken = - input.internalIntegrationSecret.length > 0 - ? encrypt(input.internalIntegrationSecret) - : existingConfig?.encryptedToken; - - if (!nextEncryptedToken) { - throw new Error( - 'A Notion internal integration secret is required when no secret is already stored.', - ); - } - - const authConfig = { - type: 'notion' as const, - encryptedToken: nextEncryptedToken, - }; - - await db - .insert(mcpConnections) - .values({ - userId: null, - mcpId: 'notion', - connectionRole: 'default', - authConfig, - enabled: true, - authStatus: 'authenticated', - }) - .onConflictDoUpdate({ - target: [ - mcpConnections.userId, - mcpConnections.mcpId, - mcpConnections.connectionRole, - ], - set: { - connectionRole: 'default', - authConfig, - accessToken: null, - refreshToken: null, - tokenExpiresAt: null, - scopes: null, - enabled: true, - authStatus: 'authenticated', - updatedAt: new Date(), - }, - }); - - await db - .insert(deploymentMcpEnablements) - .values({ - mcpId: 'notion', - enabled: true, - enabledByUserId: auth.userId, - }) - .onConflictDoUpdate({ - target: [deploymentMcpEnablements.mcpId], - set: { - enabled: true, - enabledByUserId: auth.userId, - disabledTools: null, - updatedAt: new Date(), - }, - }); + await saveAuthenticatedDeploymentMcpConnection({ + mcpId: 'notion', + userId: auth.userId, + clearOauthTokens: true, + resetDisabledTools: true, + buildAuthConfig: (existingAuthConfig) => { + const existingConfig = isMcpConnectionNotionConfig(existingAuthConfig) + ? existingAuthConfig + : null; + const encryptedToken = input.internalIntegrationSecret.length + ? encrypt(input.internalIntegrationSecret) + : existingConfig?.encryptedToken; + + if (!encryptedToken) { + throw new Error( + 'A Notion internal integration secret is required when no secret is already stored.', + ); + } - if (!existingConnection) { - captureIntegrationLifecycleEvent( - 'integration_connected', - 'notion', - auth.userId, - ); - captureIntegrationLifecycleEvent( - 'integration_enabled', - 'notion', - auth.userId, - ); - } + return { type: 'notion' as const, encryptedToken }; + }, + }); return { authStatus: 'authenticated' as const }; } @@ -1339,94 +1157,33 @@ export async function saveRipplingConnectionCommand( assertAdmin(auth); assertCuratedIntegrationsEnabled(); - const existingConnection = await db.query.mcpConnections.findFirst({ - where: and( - eq(mcpConnections.mcpId, 'rippling'), - isNull(mcpConnections.userId), - ), - columns: { authConfig: true }, - }); - const existingConfig = isMcpConnectionRipplingConfig( - existingConnection?.authConfig, - ) - ? existingConnection.authConfig - : null; - const nextEncryptedApiToken = - input.apiToken.length > 0 - ? encrypt(input.apiToken) - : existingConfig?.encryptedApiToken; - - if (!nextEncryptedApiToken) { - throw new Error( - 'A Rippling API token is required when no token is already stored.', - ); - } - - const authConfig = { - type: 'rippling' as const, - encryptedApiToken: nextEncryptedApiToken, - }; - - await validateRipplingConnection(authConfig); - - await db - .insert(mcpConnections) - .values({ - userId: null, - mcpId: 'rippling', - connectionRole: 'default', - authConfig, - enabled: true, - authStatus: 'authenticated', - }) - .onConflictDoUpdate({ - target: [ - mcpConnections.userId, - mcpConnections.mcpId, - mcpConnections.connectionRole, - ], - set: { - connectionRole: 'default', - authConfig, - accessToken: null, - refreshToken: null, - tokenExpiresAt: null, - scopes: null, - enabled: true, - authStatus: 'authenticated', - updatedAt: new Date(), - }, - }); - - await db - .insert(deploymentMcpEnablements) - .values({ - mcpId: 'rippling', - enabled: true, - enabledByUserId: auth.userId, - }) - .onConflictDoUpdate({ - target: [deploymentMcpEnablements.mcpId], - set: { - enabled: true, - enabledByUserId: auth.userId, - disabledTools: null, - updatedAt: new Date(), - }, - }); + await saveAuthenticatedDeploymentMcpConnection({ + mcpId: 'rippling', + userId: auth.userId, + clearOauthTokens: true, + resetDisabledTools: true, + buildAuthConfig: async (existingAuthConfig) => { + const existingConfig = isMcpConnectionRipplingConfig(existingAuthConfig) + ? existingAuthConfig + : null; + const encryptedApiToken = input.apiToken.length + ? encrypt(input.apiToken) + : existingConfig?.encryptedApiToken; + + if (!encryptedApiToken) { + throw new Error( + 'A Rippling API token is required when no token is already stored.', + ); + } - if (!existingConnection) { - captureIntegrationLifecycleEvent( - 'integration_connected', - 'rippling', - auth.userId, - ); - captureIntegrationLifecycleEvent( - 'integration_enabled', - 'rippling', - auth.userId, - ); - } + const authConfig = { + type: 'rippling' as const, + encryptedApiToken, + }; + await validateRipplingConnection(authConfig); + return authConfig; + }, + }); return { authStatus: 'authenticated' as const }; } @@ -1438,91 +1195,27 @@ export async function saveGranolaConnectionCommand( assertAdmin(auth); assertCuratedIntegrationsEnabled(); - const existingConnection = await db.query.mcpConnections.findFirst({ - where: and( - eq(mcpConnections.mcpId, 'granola'), - isNull(mcpConnections.userId), - ), - columns: { - authConfig: true, + await saveAuthenticatedDeploymentMcpConnection({ + mcpId: 'granola', + userId: auth.userId, + buildAuthConfig: (existingAuthConfig) => { + const existingConfig = isMcpConnectionGranolaConfig(existingAuthConfig) + ? existingAuthConfig + : null; + const encryptedApiKey = input.apiKey.length + ? encrypt(input.apiKey) + : existingConfig?.encryptedApiKey; + + if (!encryptedApiKey) { + throw new Error( + 'Granola API key is required when no Granola key is already stored.', + ); + } + + return { type: 'granola' as const, encryptedApiKey }; }, }); - const existingConfig = isMcpConnectionGranolaConfig( - existingConnection?.authConfig, - ) - ? existingConnection.authConfig - : null; - const nextEncryptedApiKey = - input.apiKey.length > 0 - ? encrypt(input.apiKey) - : existingConfig?.encryptedApiKey; - - if (!nextEncryptedApiKey) { - throw new Error( - 'Granola API key is required when no Granola key is already stored.', - ); - } - - const authConfig = { - type: 'granola' as const, - encryptedApiKey: nextEncryptedApiKey, - }; - - await db - .insert(mcpConnections) - .values({ - userId: null, - mcpId: 'granola', - connectionRole: 'default', - authConfig, - enabled: true, - authStatus: 'authenticated', - }) - .onConflictDoUpdate({ - target: [ - mcpConnections.userId, - mcpConnections.mcpId, - mcpConnections.connectionRole, - ], - set: { - connectionRole: 'default', - authConfig, - enabled: true, - authStatus: 'authenticated', - updatedAt: new Date(), - }, - }); - - await db - .insert(deploymentMcpEnablements) - .values({ - mcpId: 'granola', - enabled: true, - enabledByUserId: auth.userId, - }) - .onConflictDoUpdate({ - target: [deploymentMcpEnablements.mcpId], - set: { - enabled: true, - enabledByUserId: auth.userId, - updatedAt: new Date(), - }, - }); - - if (!existingConnection) { - captureIntegrationLifecycleEvent( - 'integration_connected', - 'granola', - auth.userId, - ); - captureIntegrationLifecycleEvent( - 'integration_enabled', - 'granola', - auth.userId, - ); - } - return { authStatus: 'authenticated' as const, }; @@ -1535,92 +1228,31 @@ export async function saveElevenLabsConnectionCommand( assertAdmin(auth); assertCuratedIntegrationsEnabled(); - const existingConnection = await db.query.mcpConnections.findFirst({ - where: and( - eq(mcpConnections.mcpId, 'elevenlabs'), - isNull(mcpConnections.userId), - ), - columns: { - authConfig: true, + await saveAuthenticatedDeploymentMcpConnection({ + mcpId: 'elevenlabs', + userId: auth.userId, + buildAuthConfig: (existingAuthConfig) => { + const existingConfig = isMcpConnectionElevenLabsConfig(existingAuthConfig) + ? existingAuthConfig + : null; + const encryptedApiKey = input.apiKey.length + ? encrypt(input.apiKey) + : existingConfig?.encryptedApiKey; + + if (!encryptedApiKey) { + throw new Error( + 'ElevenLabs API key is required when no ElevenLabs key is already stored.', + ); + } + + return { + type: 'elevenlabs' as const, + encryptedApiKey, + voiceId: input.voiceId, + }; }, }); - const existingConfig = isMcpConnectionElevenLabsConfig( - existingConnection?.authConfig, - ) - ? existingConnection.authConfig - : null; - const nextEncryptedApiKey = - input.apiKey.length > 0 - ? encrypt(input.apiKey) - : existingConfig?.encryptedApiKey; - - if (!nextEncryptedApiKey) { - throw new Error( - 'ElevenLabs API key is required when no ElevenLabs key is already stored.', - ); - } - - const authConfig = { - type: 'elevenlabs' as const, - encryptedApiKey: nextEncryptedApiKey, - voiceId: input.voiceId, - }; - - await db - .insert(mcpConnections) - .values({ - userId: null, - mcpId: 'elevenlabs', - connectionRole: 'default', - authConfig, - enabled: true, - authStatus: 'authenticated', - }) - .onConflictDoUpdate({ - target: [ - mcpConnections.userId, - mcpConnections.mcpId, - mcpConnections.connectionRole, - ], - set: { - connectionRole: 'default', - authConfig, - enabled: true, - authStatus: 'authenticated', - updatedAt: new Date(), - }, - }); - - await db - .insert(deploymentMcpEnablements) - .values({ - mcpId: 'elevenlabs', - enabled: true, - enabledByUserId: auth.userId, - }) - .onConflictDoUpdate({ - target: [deploymentMcpEnablements.mcpId], - set: { - enabled: true, - enabledByUserId: auth.userId, - updatedAt: new Date(), - }, - }); - - if (!existingConnection) { - captureIntegrationLifecycleEvent( - 'integration_connected', - 'elevenlabs', - auth.userId, - ); - captureIntegrationLifecycleEvent( - 'integration_enabled', - 'elevenlabs', - auth.userId, - ); - } - return { authStatus: 'authenticated' as const, }; @@ -1633,78 +1265,27 @@ export async function saveXConnectionCommand( assertAdmin(auth); assertCuratedIntegrationsEnabled(); - const existingConnection = await db.query.mcpConnections.findFirst({ - where: and(eq(mcpConnections.mcpId, 'x'), isNull(mcpConnections.userId)), - columns: { - authConfig: true, + await saveAuthenticatedDeploymentMcpConnection({ + mcpId: 'x', + userId: auth.userId, + buildAuthConfig: (existingAuthConfig) => { + const existingConfig = isMcpConnectionXConfig(existingAuthConfig) + ? existingAuthConfig + : null; + const encryptedBearerToken = input.bearerToken.length + ? encrypt(input.bearerToken) + : existingConfig?.encryptedBearerToken; + + if (!encryptedBearerToken) { + throw new Error( + 'X bearer token is required when no X token is already stored.', + ); + } + + return { type: 'x' as const, encryptedBearerToken }; }, }); - const existingConfig = isMcpConnectionXConfig(existingConnection?.authConfig) - ? existingConnection.authConfig - : null; - const nextEncryptedBearerToken = - input.bearerToken.length > 0 - ? encrypt(input.bearerToken) - : existingConfig?.encryptedBearerToken; - - if (!nextEncryptedBearerToken) { - throw new Error( - 'X bearer token is required when no X token is already stored.', - ); - } - - const authConfig = { - type: 'x' as const, - encryptedBearerToken: nextEncryptedBearerToken, - }; - - await db - .insert(mcpConnections) - .values({ - userId: null, - mcpId: 'x', - connectionRole: 'default', - authConfig, - enabled: true, - authStatus: 'authenticated', - }) - .onConflictDoUpdate({ - target: [ - mcpConnections.userId, - mcpConnections.mcpId, - mcpConnections.connectionRole, - ], - set: { - connectionRole: 'default', - authConfig, - enabled: true, - authStatus: 'authenticated', - updatedAt: new Date(), - }, - }); - - await db - .insert(deploymentMcpEnablements) - .values({ - mcpId: 'x', - enabled: true, - enabledByUserId: auth.userId, - }) - .onConflictDoUpdate({ - target: [deploymentMcpEnablements.mcpId], - set: { - enabled: true, - enabledByUserId: auth.userId, - updatedAt: new Date(), - }, - }); - - if (!existingConnection) { - captureIntegrationLifecycleEvent('integration_connected', 'x', auth.userId); - captureIntegrationLifecycleEvent('integration_enabled', 'x', auth.userId); - } - return { authStatus: 'authenticated' as const, }; @@ -1717,94 +1298,33 @@ export async function saveVercelConnectionCommand( assertAdmin(auth); assertCuratedIntegrationsEnabled(); - const existingConnection = await db.query.mcpConnections.findFirst({ - where: and( - eq(mcpConnections.mcpId, 'vercel'), - isNull(mcpConnections.userId), - ), - columns: { - authConfig: true, + const authConfig = await saveAuthenticatedDeploymentMcpConnection({ + mcpId: 'vercel', + userId: auth.userId, + buildAuthConfig: (existingAuthConfig) => { + const existingConfig = isMcpConnectionVercelConfig(existingAuthConfig) + ? existingAuthConfig + : null; + const encryptedAccessToken = input.accessToken.length + ? encrypt(input.accessToken) + : existingConfig?.encryptedAccessToken; + + if (!encryptedAccessToken) { + throw new Error( + 'Vercel access token is required when no Vercel token is already stored.', + ); + } + + return { + type: 'vercel' as const, + encryptedAccessToken, + ...(input.defaultTeamIdOrSlug + ? { defaultTeamIdOrSlug: input.defaultTeamIdOrSlug } + : {}), + }; }, }); - const existingConfig = isMcpConnectionVercelConfig( - existingConnection?.authConfig, - ) - ? existingConnection.authConfig - : null; - const nextEncryptedAccessToken = - input.accessToken.length > 0 - ? encrypt(input.accessToken) - : existingConfig?.encryptedAccessToken; - - if (!nextEncryptedAccessToken) { - throw new Error( - 'Vercel access token is required when no Vercel token is already stored.', - ); - } - - const authConfig = { - type: 'vercel' as const, - encryptedAccessToken: nextEncryptedAccessToken, - ...(input.defaultTeamIdOrSlug - ? { defaultTeamIdOrSlug: input.defaultTeamIdOrSlug } - : {}), - }; - - await db - .insert(mcpConnections) - .values({ - userId: null, - mcpId: 'vercel', - connectionRole: 'default', - authConfig, - enabled: true, - authStatus: 'authenticated', - }) - .onConflictDoUpdate({ - target: [ - mcpConnections.userId, - mcpConnections.mcpId, - mcpConnections.connectionRole, - ], - set: { - connectionRole: 'default', - authConfig, - enabled: true, - authStatus: 'authenticated', - updatedAt: new Date(), - }, - }); - - await db - .insert(deploymentMcpEnablements) - .values({ - mcpId: 'vercel', - enabled: true, - enabledByUserId: auth.userId, - }) - .onConflictDoUpdate({ - target: [deploymentMcpEnablements.mcpId], - set: { - enabled: true, - enabledByUserId: auth.userId, - updatedAt: new Date(), - }, - }); - - if (!existingConnection) { - captureIntegrationLifecycleEvent( - 'integration_connected', - 'vercel', - auth.userId, - ); - captureIntegrationLifecycleEvent( - 'integration_enabled', - 'vercel', - auth.userId, - ); - } - return { authStatus: 'authenticated' as const, defaultTeamIdOrSlug: authConfig.defaultTeamIdOrSlug, @@ -1818,92 +1338,31 @@ export async function saveGrafanaConnectionCommand( assertAdmin(auth); assertCuratedIntegrationsEnabled(); - const existingConnection = await db.query.mcpConnections.findFirst({ - where: and( - eq(mcpConnections.mcpId, 'grafana'), - isNull(mcpConnections.userId), - ), - columns: { - authConfig: true, + const authConfig = await saveAuthenticatedDeploymentMcpConnection({ + mcpId: 'grafana', + userId: auth.userId, + buildAuthConfig: (existingAuthConfig) => { + const existingConfig = isMcpConnectionGrafanaConfig(existingAuthConfig) + ? existingAuthConfig + : null; + const encryptedServiceAccountToken = input.serviceAccountToken.length + ? encrypt(input.serviceAccountToken) + : existingConfig?.encryptedServiceAccountToken; + + if (!encryptedServiceAccountToken) { + throw new Error( + 'Grafana service account token is required when no Grafana token is already stored.', + ); + } + + return { + type: 'grafana' as const, + baseUrl: normalizeGrafanaBaseUrl(input.baseUrl), + encryptedServiceAccountToken, + }; }, }); - const existingConfig = isMcpConnectionGrafanaConfig( - existingConnection?.authConfig, - ) - ? existingConnection.authConfig - : null; - const nextEncryptedServiceAccountToken = - input.serviceAccountToken.length > 0 - ? encrypt(input.serviceAccountToken) - : existingConfig?.encryptedServiceAccountToken; - - if (!nextEncryptedServiceAccountToken) { - throw new Error( - 'Grafana service account token is required when no Grafana token is already stored.', - ); - } - - const authConfig = { - type: 'grafana' as const, - baseUrl: normalizeGrafanaBaseUrl(input.baseUrl), - encryptedServiceAccountToken: nextEncryptedServiceAccountToken, - }; - - await db - .insert(mcpConnections) - .values({ - userId: null, - mcpId: 'grafana', - connectionRole: 'default', - authConfig, - enabled: true, - authStatus: 'authenticated', - }) - .onConflictDoUpdate({ - target: [ - mcpConnections.userId, - mcpConnections.mcpId, - mcpConnections.connectionRole, - ], - set: { - connectionRole: 'default', - authConfig, - enabled: true, - authStatus: 'authenticated', - updatedAt: new Date(), - }, - }); - - await db - .insert(deploymentMcpEnablements) - .values({ - mcpId: 'grafana', - enabled: true, - enabledByUserId: auth.userId, - }) - .onConflictDoUpdate({ - target: [deploymentMcpEnablements.mcpId], - set: { - enabled: true, - enabledByUserId: auth.userId, - updatedAt: new Date(), - }, - }); - - if (!existingConnection) { - captureIntegrationLifecycleEvent( - 'integration_connected', - 'grafana', - auth.userId, - ); - captureIntegrationLifecycleEvent( - 'integration_enabled', - 'grafana', - auth.userId, - ); - } - return { authStatus: 'authenticated' as const, baseUrl: authConfig.baseUrl, From f28eeb1b88b1975f6bf07a2be4d9df44b7f9143c Mon Sep 17 00:00:00 2001 From: "@mrubens" <2600+mrubens@users.noreply.github.com> Date: Tue, 25 Aug 2026 15:39:07 +0000 Subject: [PATCH 2/4] fix: align integration disable lock ordering --- .../commands/mcp-connections/index.test.ts | 26 +++++++++ .../trpc/commands/mcp-connections/index.ts | 56 +++++++++++-------- 2 files changed, 58 insertions(+), 24 deletions(-) diff --git a/apps/web/src/trpc/commands/mcp-connections/index.test.ts b/apps/web/src/trpc/commands/mcp-connections/index.test.ts index fb7ba02ab..016f468a6 100644 --- a/apps/web/src/trpc/commands/mcp-connections/index.test.ts +++ b/apps/web/src/trpc/commands/mcp-connections/index.test.ts @@ -116,4 +116,30 @@ describe('MCP connection lifecycle telemetry', () => { properties: { integration_id: 'asana' }, }); }); + + it('keeps concurrent credential saves and disables consistent', async () => { + for (let attempt = 0; attempt < 5; attempt += 1) { + await Promise.all([ + saveAsanaConnectionCommand(adminAuth, { + accessToken: `asana-token-${attempt}`, + }), + setDeploymentMcpEnabledCommand(adminAuth, { + mcpId: 'asana', + enabled: false, + }), + ]); + + const connection = await db.query.mcpConnections.findFirst({ + where: eq(mcpConnections.mcpId, 'asana'), + columns: { id: true }, + }); + const enablement = await db.query.deploymentMcpEnablements.findFirst({ + where: eq(deploymentMcpEnablements.mcpId, 'asana'), + columns: { enabled: true }, + }); + + expect(enablement).toBeDefined(); + expect(Boolean(connection)).toBe(enablement!.enabled); + } + }); }); diff --git a/apps/web/src/trpc/commands/mcp-connections/index.ts b/apps/web/src/trpc/commands/mcp-connections/index.ts index f2f95dece..1578bb838 100644 --- a/apps/web/src/trpc/commands/mcp-connections/index.ts +++ b/apps/web/src/trpc/commands/mcp-connections/index.ts @@ -7,6 +7,7 @@ import { and, isNull, or, + type DatabaseOrTransaction, } from '@roomote/db/server'; import { filterMcpToolDefinitions, @@ -670,32 +671,39 @@ export async function setDeploymentMcpEnabledCommand( const defaultDisabledTools = getMcpIntegrationDefaultDisabledTools(integration); - const [result] = await db - .insert(deploymentMcpEnablements) - .values({ - mcpId: input.mcpId, - enabled: input.enabled, - enabledByUserId: auth.userId, - ...(defaultDisabledTools.length > 0 - ? { disabledTools: [...defaultDisabledTools] } - : {}), - }) - .onConflictDoUpdate({ - target: [deploymentMcpEnablements.mcpId], - set: { + const persistEnablement = async (database: DatabaseOrTransaction) => { + const [result] = await database + .insert(deploymentMcpEnablements) + .values({ + mcpId: input.mcpId, enabled: input.enabled, enabledByUserId: auth.userId, - updatedAt: new Date(), - }, - }) - .returning(); + ...(defaultDisabledTools.length > 0 + ? { disabledTools: [...defaultDisabledTools] } + : {}), + }) + .onConflictDoUpdate({ + target: [deploymentMcpEnablements.mcpId], + set: { + enabled: input.enabled, + enabledByUserId: auth.userId, + updatedAt: new Date(), + }, + }) + .returning(); - // When disabling, clean up all user connections for this MCP - if (!input.enabled) { - await db - .delete(mcpConnections) - .where(eq(mcpConnections.mcpId, input.mcpId)); - } + return result!; + }; + + const result = input.enabled + ? await persistEnablement(db) + : await db.transaction(async (tx) => { + // Credential saves lock the connection before the enablement row too. + await tx + .delete(mcpConnections) + .where(eq(mcpConnections.mcpId, input.mcpId)); + return persistEnablement(tx); + }); captureIntegrationLifecycleEvent( input.enabled ? 'integration_enabled' : 'integration_disabled', @@ -703,7 +711,7 @@ export async function setDeploymentMcpEnabledCommand( auth.userId, ); - return result!; + return result; } /** From d847de7b835bf2aba8de3ef4fce82515eba1d88c Mon Sep 17 00:00:00 2001 From: "@mrubens" <2600+mrubens@users.noreply.github.com> Date: Tue, 25 Aug 2026 15:47:26 +0000 Subject: [PATCH 3/4] test: scope brain backfill idempotency assertion --- packages/db/src/lib/__tests__/brain.test.ts | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/packages/db/src/lib/__tests__/brain.test.ts b/packages/db/src/lib/__tests__/brain.test.ts index 8ed8fe580..2f43105dc 100644 --- a/packages/db/src/lib/__tests__/brain.test.ts +++ b/packages/db/src/lib/__tests__/brain.test.ts @@ -514,10 +514,9 @@ describe('backfillBrainMemoryEvents', () => { .returning(); const first = await backfillBrainMemoryEvents(db); - const second = await backfillBrainMemoryEvents(db); + await backfillBrainMemoryEvents(db); expect(first).toBeGreaterThanOrEqual(1); - expect(second).toBe(0); const completedEvents = await db .select() From b0431c6661b542c920f97fa9f7f54b4a459285ef Mon Sep 17 00:00:00 2001 From: "@mrubens" <2600+mrubens@users.noreply.github.com> Date: Tue, 25 Aug 2026 15:56:09 +0000 Subject: [PATCH 4/4] test: avoid cross-worker MCP state assertion --- .../commands/mcp-connections/index.test.ts | 26 ------------------- 1 file changed, 26 deletions(-) diff --git a/apps/web/src/trpc/commands/mcp-connections/index.test.ts b/apps/web/src/trpc/commands/mcp-connections/index.test.ts index 016f468a6..fb7ba02ab 100644 --- a/apps/web/src/trpc/commands/mcp-connections/index.test.ts +++ b/apps/web/src/trpc/commands/mcp-connections/index.test.ts @@ -116,30 +116,4 @@ describe('MCP connection lifecycle telemetry', () => { properties: { integration_id: 'asana' }, }); }); - - it('keeps concurrent credential saves and disables consistent', async () => { - for (let attempt = 0; attempt < 5; attempt += 1) { - await Promise.all([ - saveAsanaConnectionCommand(adminAuth, { - accessToken: `asana-token-${attempt}`, - }), - setDeploymentMcpEnabledCommand(adminAuth, { - mcpId: 'asana', - enabled: false, - }), - ]); - - const connection = await db.query.mcpConnections.findFirst({ - where: eq(mcpConnections.mcpId, 'asana'), - columns: { id: true }, - }); - const enablement = await db.query.deploymentMcpEnablements.findFirst({ - where: eq(deploymentMcpEnablements.mcpId, 'asana'), - columns: { enabled: true }, - }); - - expect(enablement).toBeDefined(); - expect(Boolean(connection)).toBe(enablement!.enabled); - } - }); });