fix(integrations): preserve active account authorizations

Author: waleedlatif1

apps/desktop/e2e/source-connect.spec.ts

@@ -44,8 +44,9 @@ test('source authorization returns to its desktop screen and refreshes live', as
   }
   const tickets = new Map<string, unknown>()
   const attempts = new Map<string, string>()
-  const accountAttempts = new Map<string, string>()
+  const accountAttempts = new Map<string, { session: string; mcp: boolean }>()
   let accountConnected = false
+  let mcpAccountConnected = false
   const startSessions: string[] = []
   const callbackSessions: string[] = []
   const githubAttempts = new Map<string, { session: string; completed: boolean }>()
@@ -90,6 +91,16 @@ test('source authorization returns to its desktop screen and refreshes live', as
         availableMcpConnectors: [],
         canManage: false,
         indexingAvailable: true,
+        viewerMcpAccounts: mcpAccountConnected
+          ? [
+              {
+                credentialId: 'fixture-mcp-account',
+                displayName: 'Fixture MCP account',
+                mcpServerId: 'fixture-mcp',
+                status: 'active',
+              },
+            ]
+          : [],
         viewerAccounts: accountConnected
           ? [
               {
@@ -109,15 +120,13 @@ test('source authorization returns to its desktop screen and refreshes live', as
       path === '/api/organizations/fixture-organization/connected-accounts/connect' ||
       path === '/api/users/me/organization-accounts/fixture-account/reconnect'
     ) {
-      const completionId =
-        request.method === 'POST' && path.endsWith('/connect')
-          ? (await body()).oauthCompletionId
-          : url.searchParams.get('oauthCompletionId')
+      const input = request.method === 'POST' && path.endsWith('/connect') ? await body() : null
+      const completionId = input?.oauthCompletionId ?? url.searchParams.get('oauthCompletionId')
       if (!completionId) {
         json({ error: 'Missing completion ID' }, 400)
         return
       }
-      accountAttempts.set(completionId, session)
+      accountAttempts.set(completionId, { session, mcp: Boolean(input?.mcpServerId) })
       json({
         invitationLink: `${origin}/credential-groups/enroll/fixture-account-invitation`,
         authorizationUrl: `${origin}/account-provider?completionId=${completionId}`,
@@ -126,13 +135,17 @@ test('source authorization returns to its desktop screen and refreshes live', as
     }
     if (path === '/account-callback') {
       const completionId = url.searchParams.get('completionId') ?? ''
-      if (accountAttempts.get(completionId) !== session) {
+      const attempt = accountAttempts.get(completionId)
+      if (attempt?.session !== session) {
         json({ error: 'Wrong attempt' }, 403)
         return
       }
       accountAttempts.delete(completionId)
       const denied = url.searchParams.has('error')
-      if (!denied) accountConnected = true
+      if (!denied) {
+        if (attempt.mcp) mcpAccountConnected = true
+        else accountConnected = true
+      }
       redirect(
         `/credential-groups/complete?completionId=${completionId}&organizationId=fixture-organization${denied ? '&oauth=denied' : ''}`
       )
@@ -494,6 +507,7 @@ test('source authorization returns to its desktop screen and refreshes live', as
       'web authorization preserves the origin and refreshes after an isolated provider window',
       async () => {
         accountConnected = false
+        mcpAccountConnected = false
         await web.reload()
         await web.getByLabel('Source draft').fill('Web draft retained')
         await expect(web.getByLabel('Account count')).toHaveText('0')
@@ -507,6 +521,19 @@ test('source authorization returns to its desktop screen and refreshes live', as
         expect(web.url()).toBe(`${origin}/o/fixture-organization/integrations?search=fixture`)
       }
     )
+    await check('overlapping connect and reconnect preserve the active authorization', async () => {
+      const popupReady = context.waitForEvent('page')
+      await web.getByRole('button', { name: 'Connect account', exact: true }).click()
+      const popup = await popupReady
+      await popup.getByRole('link', { name: 'Authorize account' }).waitFor()
+      const pendingAttempts = accountAttempts.size
+      await web.getByRole('button', { name: 'Reconnect account', exact: true }).click()
+      await expect(web.getByLabel('Reconnect error')).toContainText('Finish or cancel')
+      expect(accountAttempts.size).toBe(pendingAttempts)
+      await expect(web.getByLabel('Account authorization', { exact: true })).toHaveText('pending')
+      await popup.getByRole('link', { name: 'Authorize account' }).click()
+      await expect(web.getByLabel('Account authorization', { exact: true })).toHaveText('success')
+    })
     await check('web denial and cancellation leave the initiating page usable', async () => {
       const popupReady = context.waitForEvent('page')
       await web.getByRole('button', { name: 'Connect account', exact: true }).click()

apps/sim/app/o/[organizationId]/integrations/integrations.test.tsx

@@ -22,7 +22,6 @@ import { NuqsTestingAdapter } from 'nuqs/adapters/testing'
 import { createRoot, type Root } from 'react-dom/client'
 import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'
 import type { SearchSourceSummary } from '@/lib/api/contracts/knowledge/connectors'
-import type { OrganizationAccountConnectionResponse } from '@/lib/api/contracts/organization-accounts'
 import type { SearchConnector } from '@/lib/sim-search/connectors'
 
 const mocks = vi.hoisted(() => ({
@@ -287,22 +286,6 @@ function menuItem(label: string) {
   )!
 }
 
-function expectConnectionRedirect(
-  onSuccess: (response: OrganizationAccountConnectionResponse) => void,
-  authorizationUrl?: string
-) {
-  const invitationLink = 'https://sim.test/credential-groups/enroll/fixture-token'
-  const assign = vi.fn()
-  const browserWindow = window
-  vi.stubGlobal('window', { location: { assign } })
-  try {
-    onSuccess({ invitationLink, ...(authorizationUrl ? { authorizationUrl } : {}) })
-    expect(assign).toHaveBeenCalledExactlyOnceWith(authorizationUrl ?? invitationLink)
-  } finally {
-    vi.stubGlobal('window', browserWindow)
-  }
-}
-
 describe('GitHub member account inventory', () => {
   const githubAccount = {
     credentialId: 'github-account',
@@ -346,10 +329,7 @@ describe('GitHub member account inventory', () => {
     })
   })
 
-  it.each([
-    undefined,
-    'https://sim.test/api/credential-groups/enroll/fixture-token/oauth/github-option?returnTo=search',
-  ])('connects once through the account operation with compatible redirect %s', async (url) => {
+  it('connects once through the account operation', async () => {
     await render()
     expect(buttons('Connect')).toHaveLength(1)
     expect(container.textContent).toContain('Connect once')
@@ -358,7 +338,6 @@ describe('GitHub member account inventory', () => {
       { organizationId: scope.organizationId, optionId: 'github-option' },
       expect.any(Object)
     )
-    expectConnectionRedirect(mocks.connectOrganizationAccount.mock.calls[0][1].onSuccess, url)
     expect(mockUseSearchSources).not.toHaveBeenCalled()
     expect(mocks.connect).not.toHaveBeenCalled()
     expect(mocks.connectSearchSource).not.toHaveBeenCalled()
@@ -430,10 +409,7 @@ describe('GitHub member account inventory', () => {
     }
   )
 
-  it.each([
-    undefined,
-    'https://sim.test/api/credential-groups/enroll/fixture-token/oauth/github-option?returnTo=accounts',
-  ])('allows personal reauthorization while Search is disabled with redirect %s', async (url) => {
+  it('allows personal reauthorization while Search is disabled', async () => {
     mockUseSearchSourceOverview.mockReturnValue({ data: { providers: [] }, isPending: false })
     mocks.integrations.mockReturnValue({
       data: [{ connectorType: 'github', approved: false }],
@@ -457,7 +433,6 @@ describe('GitHub member account inventory', () => {
       'github-account',
       expect.any(Object)
     )
-    expectConnectionRedirect(mocks.reconnectOrganizationAccount.mock.calls[0][1].onSuccess, url)
     expect(mocks.connect).not.toHaveBeenCalled()
   })
 

apps/sim/hooks/queries/organization-accounts.ts

@@ -63,7 +63,11 @@ function useAccountConnectionMutation<Variables>(
   const pending = useRef<AbortController | null>(null)
   useEffect(() => () => pending.current?.abort(), [])
   return useMutation({
+    mutationKey: organizationAccountsKeys.connection(),
     mutationFn: async (variables: Variables) => {
+      if (client.isMutating({ mutationKey: organizationAccountsKeys.connection() }) > 1) {
+        throw new Error('Finish or cancel your current account connection before starting another.')
+      }
       pending.current?.abort()
       const controller = new AbortController()
       pending.current = controller
@@ -131,6 +135,7 @@ export function useDisconnectPersonalOrganizationAccount(organizationId: string)
 
 export const organizationAccountsKeys = {
   all: ['organization-accounts'] as const,
+  connection: () => [...organizationAccountsKeys.all, 'connection'] as const,
   workspaces: () => [...organizationAccountsKeys.all, 'workspace'] as const,
   workspace: (workspaceId?: string) =>
     [...organizationAccountsKeys.workspaces(), workspaceId ?? ''] as const,

apps/sim/scripts/fixtures/desktop-source-connect.tsx

@@ -2,6 +2,7 @@ import { StrictMode, useEffect, useRef, useState } from 'react'
 import { ToastProvider } from '@sim/emcn'
 import { QueryClient, QueryClientProvider } from '@tanstack/react-query'
 import { createRoot } from 'react-dom/client'
+import { isCredentialGroupOAuthFailure } from '@/lib/credential-groups/oauth-completion'
 import { startDesktopSourceBrowser } from '@/lib/desktop/source-browser'
 import { CredentialGroupCompletionHandoff } from '@/app/credential-groups/complete/completion-handoff'
 import { SlackCompletion } from '@/app/credential-groups/slack-complete/slack-completion'
@@ -63,7 +64,11 @@ function SourceConnectFixture() {
       <output aria-label='Account authorization'>{accountConnection.status}</output>
       <output aria-label='Account error'>{accountConnection.error?.message}</output>
       <output aria-label='Reconnect status'>{reconnect.status}</output>
-      <output aria-label='Account count'>{accounts.data?.viewerAccounts?.length ?? 0}</output>
+      <output aria-label='Reconnect error'>{reconnect.error?.message}</output>
+      <output aria-label='Account count'>
+        {(accounts.data?.viewerAccounts?.length ?? 0) +
+          (accounts.data?.viewerMcpAccounts?.length ?? 0)}
+      </output>
       <input aria-label='Source draft' defaultValue='Unsubmitted source name' />
       <button
         disabled={connection.isPending}
@@ -118,6 +123,8 @@ function BrowserLauncher() {
 }
 
 const params = new URLSearchParams(location.search)
+const oauth = params.get('oauth')
+const failure = oauth === null ? undefined : isCredentialGroupOAuthFailure(oauth) ? oauth : 'failed'
 const content =
   location.pathname === '/credential-groups/enroll/fixture-invitation' ? (
     params.has('connected') ? (
@@ -130,7 +137,7 @@ const content =
   ) : location.pathname === '/credential-groups/complete' ? (
     <CredentialGroupCompletionHandoff
       completionId={params.get('completionId')!}
-      failure={params.get('oauth') === 'denied' ? 'denied' : undefined}
+      failure={failure}
       returnHref={params.has('organizationId') ? '/o/fixture-organization/integrations' : undefined}
     />
   ) : location.pathname === '/desktop/connect' ? (