fix(desktop): preserve enrollment mode and completion state

Author: waleedlatif1

.github/workflows/desktop-e2e.yml

@@ -17,6 +17,9 @@ on:
       - 'apps/sim/app/desktop/connect/**'
       - 'apps/sim/app/credential-groups/**'
       - 'apps/sim/hooks/queries/slack-search.ts'
+      - 'apps/sim/hooks/use-github-installation-setup.ts'
+      - 'apps/sim/app/o/**/integrations/indexed/use-member-enrollment.ts'
+      - 'apps/sim/lib/api/contracts/desktop-source-connect.ts'
       - 'apps/sim/scripts/fixtures/desktop-source-connect.tsx'
       - 'apps/sim/app/workspace/**/browser-session/**'
       - 'apps/sim/app/_styles/**'

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

@@ -113,6 +113,19 @@ test('source authorization returns to its desktop screen and refreshes live', as
       redirect(`/credential-groups/slack-complete?state=${state}&ok=${ok}`)
       return
     }
+    if (
+      path ===
+      '/api/knowledge/00000000-0000-4000-8000-000000000001/connectors/fixture-connector/enroll'
+    ) {
+      if (url.searchParams.has('oauthCompletionId'))
+        json({ error: 'Direct account connection requires a Search source' }, 400)
+      else
+        json({
+          success: true,
+          data: { url: `${origin}/credential-groups/enroll/fixture-invitation` },
+        })
+      return
+    }
     if (path === '/api/knowledge/github/setup') {
       if (request.method === 'POST') {
         const { setupId } = await body()
@@ -367,6 +380,15 @@ test('source authorization returns to its desktop screen and refreshes live', as
         expect(page.url()).toBe(`${origin}/home`)
       }
     )
+    await check('ordinary knowledge-base enrollment preserves its invitation step', async () => {
+      await page.getByRole('button', { name: 'Connect invited source' }).click()
+      await expect.poll(async () => (await opened()).length).toBe(7)
+      await external.goto((await opened())[6])
+      await external.getByRole('link', { name: 'Authorize invited source' }).click()
+      await expect(page.getByLabel('Enrollment pending')).toHaveText('false')
+      await expect(page.getByLabel('Enrollment error')).toHaveText('')
+      expect(page.url()).toBe(`${origin}/home`)
+    })
     await page.screenshot({ path: test.info().outputPath('source-connect-desktop.png') })
   } finally {
     mkdirSync(dirname(reportPath), { recursive: true })

apps/sim/app/api/knowledge/slack/oauth/callback/route.ts

@@ -1,3 +1,5 @@
+import { createLogger } from '@sim/logger'
+import { describeError } from '@sim/utils/errors'
 import { NextResponse } from 'next/server'
 import { slackSearchOAuthCallbackContract } from '@/lib/api/contracts/knowledge/slack'
 import { parseRequest } from '@/lib/api/server'
@@ -15,6 +17,8 @@ import { completeSlackSearchSetup } from '@/lib/knowledge/application/slack-sear
 import { slackSearchInstallPath } from '@/lib/slack-search/install-link'
 import { authenticateSlackPublicInstallation } from '@/lib/slack-search/public-install-auth'
 
+const logger = createLogger('SlackSearchOAuthCallback')
+
 /** OAuth is a redirect protocol; protected configuration remains in the application use case. */
 export const GET = withRouteHandler(async (request) => {
   let callbackState: string | undefined
@@ -64,6 +68,12 @@ export const GET = withRouteHandler(async (request) => {
     })
   } catch (error) {
     if (callbackState) {
+      const projected = internalOrchestrationErrorPolicy.project(error)
+      if (
+        !(error instanceof InternalUnauthenticatedError) &&
+        (!projected || projected.status >= 500)
+      )
+        logger.error('Slack authorization callback failed', { error: describeError(error) })
       const url = new URL('/credential-groups/slack-complete', getBaseUrl())
       url.searchParams.set('state', callbackState)
       url.searchParams.set('ok', 'false')

apps/sim/app/o/[organizationId]/integrations/indexed/use-member-enrollment.ts

@@ -108,7 +108,10 @@ export function useMemberEnrollment({
       nativeAbort.current = controller
       return connectDesktopSource(request, controller.signal)
     },
-    onSettled: () => refreshMemberships(),
+    onSettled: () =>
+      Promise.all(
+        membershipQueryKeys.map((queryKey) => queryClient.invalidateQueries({ queryKey }))
+      ),
     onError: (error) => onConnectionError?.(error.message),
     onSuccess: () => setSetupConnector(null),
   })
@@ -286,7 +289,7 @@ export function useMemberEnrollment({
       nativeConnection.mutate({
         kind: 'member-enrollment',
         params: { id: knowledgeBaseId, connectorId },
-        completionId: generateId(),
+        ...(directOAuth ? { completionId: generateId() } : {}),
       })
       return
     }
@@ -321,7 +324,7 @@ export function useMemberEnrollment({
       nativeConnection.mutate({
         kind: 'search-source',
         body: { ...resourceScopeFields(scope), connectorType, sourceConfig },
-        completionId: generateId(),
+        ...(directOAuth ? { completionId: generateId() } : {}),
       })
       return
     }

apps/sim/components/integrations/slack-search-setup-wizard.tsx

@@ -156,7 +156,7 @@ export function SlackSearchSetupWizard({
     return (
       <ChipModal
         open
-        dismissDisabled={busy && !isDesktopApp()}
+        dismissDisabled={connect.isPending || (oauth.isPending && !isDesktopApp())}
         onOpenChange={(open) => {
           if (!open) close()
         }}
@@ -212,7 +212,7 @@ export function SlackSearchSetupWizard({
   return (
     <ChipModal
       open
-      dismissDisabled={busy && !isDesktopApp()}
+      dismissDisabled={connect.isPending || (oauth.isPending && !isDesktopApp())}
       onOpenChange={(open) => {
         if (!open) close()
       }}

apps/sim/hooks/queries/kb/connectors.test.ts

@@ -2,11 +2,13 @@ import {
   apiClientRequestMock,
   apiClientRequestMockFns,
 } from '@sim/testing/mocks/api-client-request.mock'
+import { emcnMock } from '@sim/testing/mocks/emcn.mock'
 import { reactQueryMock, reactQueryMockFns } from '@sim/testing/mocks/react-query.mock'
 import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'
 import { searchSourceKeys } from '@/hooks/queries/utils/search-source-keys'
 
 vi.mock('react', () => ({ useEffect: vi.fn() }))
+vi.mock('@sim/emcn', () => emcnMock)
 
 vi.mock('@tanstack/react-query', () => reactQueryMock)
 

apps/sim/hooks/queries/personal-search-integrations.ts

@@ -45,7 +45,7 @@ export function useConnectPersonalSearchIntegration() {
         await connectDesktopSource({ kind: 'personal-search', body }, signal)
         return null
       }
-      return (await requestJson(connectPersonalSearchIntegrationContract, { body })).data
+      return (await requestJson(connectPersonalSearchIntegrationContract, { body, signal })).data
     },
     onSettled: (_data, _error, body) =>
       Promise.all([

apps/sim/lib/api/contracts/desktop-source-connect.ts

@@ -25,12 +25,12 @@ export const desktopSourceRequestSchema = z.discriminatedUnion('kind', [
   z.object({
     kind: z.literal('member-enrollment'),
     params: knowledgeConnectorParamsSchema,
-    completionId: z.string().uuid(),
+    completionId: z.string().uuid().optional(),
   }),
   z.object({
     kind: z.literal('search-source'),
     body: connectSimSearchConnectorBodySchema,
-    completionId: z.string().uuid(),
+    completionId: z.string().uuid().optional(),
   }),
   z.object({
     kind: z.literal('slack-managed-users'),

apps/sim/lib/desktop/source-browser.ts

@@ -106,13 +106,23 @@ async function startRequest(
         params: request.params,
         query: { oauthCompletionId: request.completionId },
       })
-      return { url: result.data.url, match: { kind: 'completion', id: request.completionId } }
+      return {
+        url: result.data.url,
+        match: request.completionId
+          ? { kind: 'completion', id: request.completionId }
+          : enrollmentMatch(result.data.url),
+      }
     }
     case 'search-source': {
       const result = await requestJson(connectSimSearchConnectorContract, {
         body: { ...request.body, oauthCompletionId: request.completionId },
       })
-      return { url: result.data.url, match: { kind: 'completion', id: request.completionId } }
+      return {
+        url: result.data.url,
+        match: request.completionId
+          ? { kind: 'completion', id: request.completionId }
+          : enrollmentMatch(result.data.url),
+      }
     }
     case 'slack-managed-users': {
       const { owner, body, credentialGroupId } = request

apps/sim/lib/desktop/source-request.integration.ts

@@ -3,9 +3,11 @@ import { sleep } from '@sim/utils/helpers'
 import { generateId, generateShortId } from '@sim/utils/id'
 import { afterAll, describe, expect, it, vi } from 'vitest'
 
-await vi.hoisted(async () => {
+const { redisUrl } = await vi.hoisted(async () => {
   const { readTestRedisUrl } = await import('@sim/db/testing/test-infrastructure')
-  process.env.REDIS_URL = readTestRedisUrl()
+  const redisUrl = readTestRedisUrl()
+  if (redisUrl) process.env.REDIS_URL = redisUrl
+  return { redisUrl }
 })
 
 import { closeRedisConnection, getRedisClient } from '@/lib/core/config/redis'
@@ -18,7 +20,7 @@ import {
 } from '@/lib/desktop/application/source-requests'
 
 /** Real Redis proves cross-session ownership and atomic consumption, without provider credentials. */
-describe('desktop source request transport', () => {
+describe.skipIf(!redisUrl)('desktop source request transport', () => {
   it('allows the same user in another session, without letting another user consume the request', async () => {
     const userId = generateId()
     const created = await createDesktopSourceRequest.execute({