fix(search): validate Atlassian keys beyond discovery limits

Author: waleedlatif1

apps/sim/lib/knowledge/application/personal-source-setup.test.ts

@@ -1,5 +1,5 @@
 /** @vitest-environment node */
-import { beforeEach, describe, expect, it, vi } from 'vitest'
+import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'
 
 const mocks = vi.hoisted(() => ({
   authorizeOperation: vi.fn(),
@@ -69,6 +69,12 @@ import {
   listPersonalSourceSetupAccounts,
   personalSourceSetup,
 } from '@/lib/knowledge/application/personal-source-setup'
+import { MAX_SELECTOR_PAGES } from '@/lib/selectors/limits'
+import type { SelectorRequest } from '@/lib/selectors/types'
+
+interface ValidationSelectorCall {
+  input: { request: SelectorRequest; signal: AbortSignal }
+}
 
 const principal = { kind: 'session', userId: 'member-1', sessionId: 'session-1' } as const
 const owner = { organizationId: 'organization-1', connectorType: 'jira' } as const
@@ -84,6 +90,10 @@ const runConnect = (changes = {}) =>
   personalSourceSetup.execute({ principal, input: { ...connect, keys: ['PROJECT'], ...changes } })
 
 describe('personal source setup', () => {
+  afterEach(() => {
+    vi.restoreAllMocks()
+  })
+
   beforeEach(() => {
     vi.resetAllMocks()
     mocks.authorize.mockResolvedValue(principal.userId)
@@ -245,15 +255,154 @@ describe('personal source setup', () => {
     { kind: 'list', items: [], truncated: true, nextCursor: '50' },
     { kind: 'list', items: [] },
   ])('rejects unavailable manually entered keys before source creation %#', async (page) => {
-    mocks.selector.mockResolvedValue(page)
+    mocks.selector.mockResolvedValueOnce(page).mockResolvedValue({ kind: 'detail', item: null })
     await expect(runConnect({ keys: ['MISSING'] })).rejects.toThrow('could not be found')
     expect(mocks.configure).not.toHaveBeenCalled()
   })
 
-  it('rejects a repeated pagination cursor without looping or creating a source', async () => {
-    mocks.selector.mockResolvedValue({ kind: 'list', items: [], nextCursor: '50' })
+  it.each(['jira', 'confluence'] as const)(
+    'verifies a manually entered %s key outside the available listing through the same authorized selector',
+    async (connectorType) => {
+      mocks.selector
+        .mockResolvedValueOnce({ kind: 'list', items: [] })
+        .mockResolvedValueOnce({ kind: 'detail', item: { id: 'PROJECT', label: 'Project' } })
+      await expect(runConnect({ connectorType })).resolves.toMatchObject({ kind: 'connected' })
+      expect(mocks.selector).toHaveBeenLastCalledWith({
+        principal,
+        request: undefined,
+        input: {
+          selectorKey: connectorType === 'jira' ? 'jira.projectKeys' : 'confluence.spaces',
+          scope: { kind: 'organization', organizationId: owner.organizationId },
+          context: { oauthCredential: credential.credentialId, domain: credential.domain },
+          personalSearchSetup: connectorType,
+          signal: expect.any(AbortSignal),
+          request: { kind: 'detail', id: 'PROJECT' },
+        },
+      })
+      expect(mocks.ownAccount).toHaveBeenCalledTimes(2)
+      expect(mocks.oauth).not.toHaveBeenCalled()
+    }
+  )
+
+  it.each([
+    { name: 'truncated results', pages: 1, truncated: true, nextCursor: '50' },
+    { name: 'a repeated cursor', pages: 2, nextCursor: '50' },
+    { name: 'the page limit', pages: MAX_SELECTOR_PAGES },
+  ])('resolves remaining keys directly after $name', async ({ pages, truncated, nextCursor }) => {
+    let listed = 0
+    mocks.selector.mockImplementation(({ input }: ValidationSelectorCall) => {
+      if (input.request.kind === 'detail') {
+        return { kind: 'detail', item: { id: input.request.id, label: 'Project' } }
+      }
+      listed++
+      return { kind: 'list', items: [], nextCursor: nextCursor ?? String(listed), truncated }
+    })
+    await expect(runConnect()).resolves.toMatchObject({ kind: 'connected' })
+    expect(listed).toBe(pages)
+    expect(mocks.selector).toHaveBeenCalledTimes(pages + 1)
+    expect(mocks.oauth).not.toHaveBeenCalled()
+  })
+
+  it('gives direct validation a fresh deadline when listing times out', async () => {
+    const listing = new AbortController()
+    const details = new AbortController()
+    vi.spyOn(AbortSignal, 'timeout')
+      .mockReturnValueOnce(listing.signal)
+      .mockReturnValueOnce(details.signal)
+    mocks.selector
+      .mockImplementationOnce(({ input }: ValidationSelectorCall) => {
+        listing.abort(new DOMException('Listing timed out', 'TimeoutError'))
+        input.signal.throwIfAborted()
+      })
+      .mockImplementationOnce(({ input }: ValidationSelectorCall) => {
+        expect(input.signal.aborted).toBe(false)
+        expect(input.request).toEqual({ kind: 'detail', id: 'PROJECT' })
+        return { kind: 'detail', item: { id: 'PROJECT', label: 'Project' } }
+      })
+    await expect(runConnect()).resolves.toMatchObject({ kind: 'connected' })
+    expect(AbortSignal.timeout).toHaveBeenCalledTimes(2)
+    expect(mocks.oauth).not.toHaveBeenCalled()
+  })
+
+  it('fails before mutation when direct validation also exceeds its deadline', async () => {
+    const listing = new AbortController()
+    const details = new AbortController()
+    vi.spyOn(AbortSignal, 'timeout')
+      .mockReturnValueOnce(listing.signal)
+      .mockReturnValueOnce(details.signal)
+    mocks.selector
+      .mockResolvedValueOnce({ kind: 'list', items: [] })
+      .mockImplementationOnce(({ input }: ValidationSelectorCall) => {
+        details.abort(new DOMException('Validation timed out', 'TimeoutError'))
+        input.signal.throwIfAborted()
+      })
+    await expect(runConnect()).rejects.toThrow('took too long')
+    expect(mocks.configure).not.toHaveBeenCalled()
+  })
+
+  it('does not retry direct validation after the caller cancels listing', async () => {
+    const caller = new AbortController()
+    mocks.selector.mockImplementationOnce(({ input }: ValidationSelectorCall) => {
+      caller.abort(new Error('Setup cancelled'))
+      input.signal.throwIfAborted()
+    })
+    await expect(
+      personalSourceSetup.execute({
+        principal,
+        input: { ...connect, keys: ['PROJECT'] },
+        request: { headers: new Headers(), signal: caller.signal },
+      })
+    ).rejects.toThrow('Setup cancelled')
+    expect(mocks.selector).toHaveBeenCalledTimes(1)
+    expect(mocks.configure).not.toHaveBeenCalled()
+  })
+
+  it('rejects a detail response for a different key', async () => {
+    mocks.selector
+      .mockResolvedValueOnce({ kind: 'list', items: [] })
+      .mockResolvedValueOnce({ kind: 'detail', item: { id: 'OTHER', label: 'Other project' } })
     await expect(runConnect()).rejects.toThrow('could not be found')
-    expect(mocks.selector).toHaveBeenCalledTimes(2)
+    expect(mocks.configure).not.toHaveBeenCalled()
+  })
+
+  it('bounds direct validation concurrency and validates each unresolved key only once', async () => {
+    let release = () => {}
+    let allStarted = () => {}
+    const gate = new Promise<void>((resolve) => {
+      release = resolve
+    })
+    const started = new Promise<void>((resolve) => {
+      allStarted = resolve
+    })
+    let active = 0
+    let maximumActive = 0
+    mocks.selector.mockImplementation(async ({ input }: ValidationSelectorCall) => {
+      if (input.request.kind === 'list') return { kind: 'list', items: [] }
+      active++
+      maximumActive = Math.max(maximumActive, active)
+      if (active === 5) allStarted()
+      await gate
+      active--
+      return { kind: 'detail', item: { id: input.request.id, label: input.request.id } }
+    })
+    const keys = Array.from({ length: 12 }, (_, index) => `PROJECT${index}`)
+    const connection = runConnect({ keys: [...keys, ...keys] })
+    await started
+    try {
+      expect(mocks.selector).toHaveBeenCalledTimes(6)
+      expect(mocks.configure).not.toHaveBeenCalled()
+    } finally {
+      release()
+    }
+    await expect(connection).resolves.toMatchObject({ kind: 'connected' })
+    expect(maximumActive).toBe(5)
+    expect(mocks.selector).toHaveBeenCalledTimes(keys.length + 1)
+  })
+
+  it('propagates listing failures without starting direct validation', async () => {
+    mocks.selector.mockRejectedValueOnce(new Error('Account revoked'))
+    await expect(runConnect()).rejects.toThrow('Account revoked')
+    expect(mocks.selector).toHaveBeenCalledTimes(1)
     expect(mocks.configure).not.toHaveBeenCalled()
   })
 
@@ -274,6 +423,37 @@ describe('personal source setup', () => {
     expect(mocks.configure).not.toHaveBeenCalled()
   })
 
+  it.each(['binding', 'ownership'] as const)(
+    'stops before source creation if the caller cancels during the final %s check',
+    async (phase) => {
+      const caller = new AbortController()
+      if (phase === 'binding') {
+        mocks.binding.mockImplementationOnce(() => {
+          caller.abort(new Error('Setup cancelled'))
+          return {
+            organizationId: owner.organizationId,
+            credentialGroupId: 'group-1',
+            credentialGroupOptionId: 'option-1',
+          }
+        })
+      } else {
+        mocks.ownAccount.mockResolvedValueOnce(account).mockImplementationOnce(() => {
+          caller.abort(new Error('Setup cancelled'))
+          return account
+        })
+      }
+      await expect(
+        personalSourceSetup.execute({
+          principal,
+          input: { ...connect, keys: ['PROJECT'] },
+          request: { headers: new Headers(), signal: caller.signal },
+        })
+      ).rejects.toThrow('Setup cancelled')
+      expect(mocks.configure).not.toHaveBeenCalled()
+      expect(mocks.dispatch).not.toHaveBeenCalled()
+    }
+  )
+
   it('rejects a current operation authorization denial before any setup effects', async () => {
     mocks.authorizeOperation.mockRejectedValue(new Error('Membership ended'))
     await expect(runConnect()).rejects.toThrow('Membership ended')

apps/sim/lib/knowledge/application/personal-source-setup.ts

@@ -2,6 +2,7 @@ import { createLogger } from '@sim/logger'
 import { getErrorMessage } from '@sim/utils/errors'
 import { normalizeAtlassianSiteUrl } from '@/lib/atlassian/discovery'
 import { OrchestrationError } from '@/lib/core/orchestration/types'
+import { mapWithConcurrency } from '@/lib/core/utils/concurrency'
 import {
   loadManagedCredentialGroupBinding,
   loadScopedAccountsCredentialListContext,
@@ -29,6 +30,8 @@ import { MAX_PERSONAL_SOURCE_SETUP_KEYS } from '@/lib/sim-search/personal-source
 import { CONNECTOR_META_REGISTRY } from '@/connectors/registry'
 
 const logger = createLogger('PersonalSourceSetup')
+const VALIDATION_PHASE_TIMEOUT_MS = 30_000
+const DETAIL_VALIDATION_CONCURRENCY = 5
 
 interface PersonalSourceSetupOwner {
   organizationId: string
@@ -176,12 +179,13 @@ export const personalSourceSetup = defineAuthorizedKnowledgeUseCase({
     const keys = [...new Set(input.keys.map((key) => key.trim()))]
     const remaining = new Set(keys)
     const cursors = new Set<string>()
-    const timeout = AbortSignal.timeout(30_000)
+    const timeout = AbortSignal.timeout(VALIDATION_PHASE_TIMEOUT_MS)
     const signal = request?.signal ? AbortSignal.any([request.signal, timeout]) : timeout
     let cursor: string | undefined
     for (let page = 0; page < MAX_SELECTOR_PAGES; page++) {
       let result: SelectorExecutionResult
       try {
+        signal.throwIfAborted()
         result = await executeSelector.execute({
           principal,
           request,
@@ -191,13 +195,9 @@ export const personalSourceSetup = defineAuthorizedKnowledgeUseCase({
             request: { kind: 'list', ...(cursor ? { cursor } : {}) },
           },
         })
+        signal.throwIfAborted()
       } catch (error) {
-        if (timeout.aborted && !request?.signal?.aborted) {
-          throw new OrchestrationError(
-            'validation',
-            'Checking the selected projects or spaces took too long. Try fewer selections.'
-          )
-        }
+        if (timeout.aborted && !request?.signal?.aborted) break
         throw error
       }
       if (result.kind !== 'list') throw new Error('Source discovery returned an unexpected result')
@@ -207,11 +207,45 @@ export const personalSourceSetup = defineAuthorizedKnowledgeUseCase({
       cursor = result.nextCursor
       cursors.add(cursor)
     }
+    request?.signal?.throwIfAborted()
     if (remaining.size > 0) {
-      throw new OrchestrationError(
-        'validation',
-        'Some selected projects or spaces could not be found with this account. Refresh the choices and try again.'
-      )
+      const detailTimeout = AbortSignal.timeout(VALIDATION_PHASE_TIMEOUT_MS)
+      const details = new AbortController()
+      const detailSignal = AbortSignal.any([
+        detailTimeout,
+        details.signal,
+        ...(request?.signal ? [request.signal] : []),
+      ])
+      try {
+        await mapWithConcurrency([...remaining], DETAIL_VALIDATION_CONCURRENCY, async (key) => {
+          detailSignal.throwIfAborted()
+          const result = await executeSelector.execute({
+            principal,
+            request,
+            input: {
+              ...selectorInput,
+              signal: detailSignal,
+              request: { kind: 'detail', id: key },
+            },
+          })
+          detailSignal.throwIfAborted()
+          if (result.kind !== 'detail' || result.item?.id !== key) {
+            throw new OrchestrationError(
+              'validation',
+              'Some selected projects or spaces could not be found with this account. Refresh the choices and try again.'
+            )
+          }
+        })
+      } catch (error) {
+        details.abort(error)
+        if (detailTimeout.aborted && !request?.signal?.aborted) {
+          throw new OrchestrationError(
+            'validation',
+            'Checking the selected projects or spaces took too long. Try fewer selections.'
+          )
+        }
+        throw error
+      }
     }
     const [binding, group] = await Promise.all([
       loadManagedCredentialGroupBinding(input.credentialId),
@@ -232,6 +266,7 @@ export const personalSourceSetup = defineAuthorizedKnowledgeUseCase({
       )
     }
     await authorizePersonalSearchSetupCredential(principal, input)
+    request?.signal?.throwIfAborted()
     const result = await configureSimSearchConnector.execute({
       principal,
       request,