improvement(search): overlap live search stages and report every unsearched source

apps/sim/lib/core/utils/concurrency.ts

@@ -56,6 +56,27 @@ export function mapSettledWithConcurrency<T, R>(
   })
 }
 
+/**
+ * A gate that runs at most `limit` of the tasks passed to it at once, starting the rest in call
+ * order as slots free. Bounds one phase of otherwise concurrent work, such as the part that holds
+ * database connections, without bounding the rest of it.
+ */
+export function createConcurrencyLimit(limit: number) {
+  let active = 0
+  const waiting: (() => void)[] = []
+  return async <R>(task: () => Promise<R>): Promise<R> => {
+    if (active >= limit) await new Promise<void>((resolve) => waiting.push(resolve))
+    else active++
+    try {
+      return await task()
+    } finally {
+      const next = waiting.shift()
+      if (next) next()
+      else active--
+    }
+  }
+}
+
 /** Default bound for per-row object-storage materialization fan-out. */
 export const MATERIALIZE_CONCURRENCY = 20
 

apps/sim/lib/sim-search/live/account-session.ts

@@ -11,14 +11,15 @@ import {
   createNativeClient,
   NATIVE_SEARCH_REQUEST_BUDGET,
   NativeSearchError,
+  withRequestSignal,
 } from '@/lib/sim-search/live/http'
 import { readHubSpotMcp, searchHubSpotMcp } from '@/lib/sim-search/live/hubspot-mcp'
 import { readLucidMcp, searchLucidMcp } from '@/lib/sim-search/live/lucid-mcp'
 import { createManagedSearchMcpClient } from '@/lib/sim-search/live/managed-mcp'
 import { isManagedSearchMcpProvider } from '@/lib/sim-search/live/managed-mcp-config'
 import { readNotionMcp, searchNotionMcp } from '@/lib/sim-search/live/notion-mcp'
 import { createPolicyVerifier } from '@/lib/sim-search/live/policy'
-import type { LiveSearchPolicy } from '@/lib/sim-search/live/policy-schema'
+import { type LiveSearchPolicy, requiresScopedRetrieval } from '@/lib/sim-search/live/policy-schema'
 import { livePolicyFor, loadLiveSearchPolicies } from '@/lib/sim-search/live/policy-store'
 import { LIVE_SEARCH_PROVIDER_CATALOG } from '@/lib/sim-search/live/provider-catalog'
 import { isSearchProviderEnabled } from '@/lib/sim-search/live/provider-rollout'
@@ -51,8 +52,16 @@ export interface LiveAccountSession {
   /** True only when the document is inside the source boundary and the member may read it. */
   verify(document: Reference): Promise<boolean>
   read(reference: Reference, filters?: WorkspaceSearchFilters): Promise<NativeDocument>
-  /** Full text of a verified search result whose provider returned only an excerpt or metadata. */
-  previewText(document: NativeDocument): Promise<string | undefined>
+  /**
+   * Full text of a search result whose provider returned only an excerpt or metadata; `signal`
+   * cancels the read when its preview is no longer wanted.
+   */
+  previewText(document: NativeDocument, signal: AbortSignal): Promise<string | undefined>
+  /**
+   * True when `verify` spends requests from the member budget `previewText` also draws on, so a
+   * preview read before verification finishes could leave a later result unverified.
+   */
+  verifySharesPreviewBudget: boolean
   /** Checks a fetched document against the source settings saved now, not at session start. */
   verifyCurrent(document: NativeDocument): Promise<boolean>
 }
@@ -174,14 +183,15 @@ export async function openLiveAccountSession(
       signal,
       pool: input.pool,
     })
-    if (service) return service
+    if (service) return { ...service, spendsMemberBudget: service.spendsMemberBudget === true }
     const verifyPolicy = createPolicyVerifier(provider, memberPolicy, client, origin, mcp, {
       fresh,
     })
     return {
       policy: memberPolicy,
       partial: false,
       scopeSearch: undefined,
+      spendsMemberBudget: requiresScopedRetrieval(provider, memberPolicy),
       verify: (document: Reference) => verifyPolicy(document, document.accessMetadata),
     }
   }
@@ -196,6 +206,7 @@ export async function openLiveAccountSession(
     },
     policy: boundary.policy,
     servicePartial: boundary.partial,
+    verifySharesPreviewBudget: boundary.spendsMemberBudget,
     async search(search) {
       const searchOnce = async (request: NativeSearchInput): Promise<NativePage> => {
         const scoped = boundary.scopeSearch ? boundary.scopeSearch(request) : request
@@ -226,9 +237,9 @@ export async function openLiveAccountSession(
         })
       return readMcp(reference)
     },
-    async previewText(document) {
+    async previewText(document, previewSignal) {
       if (admin || !client) return undefined
-      return previewNativeText(provider, client, document)
+      return previewNativeText(provider, withRequestSignal(client, previewSignal), document)
     },
     async verifyCurrent(document) {
       if (!(await isSearchProviderEnabled(provider, resourceScopeFromOwner(owner)))) return false

apps/sim/lib/sim-search/live/application.test.ts

@@ -1,5 +1,5 @@
 import { ErrorCode, McpError } from '@modelcontextprotocol/sdk/types.js'
-import { dbChainMockFns, resetDbChainMock } from '@sim/testing'
+import { createDeferred, dbChainMockFns, resetDbChainMock } from '@sim/testing'
 import { createSessionPrincipal } from '@sim/testing/factories/principal.factory'
 import { setEnv } from '@sim/testing/mocks/env.mock'
 import {
@@ -740,6 +740,76 @@ describe('authorized live retrieval', () => {
       })
       expect(mocks.preview).not.toHaveBeenCalled()
     })
+
+    it('reads a permitted result’s preview while others verify, and never a denied one’s', async () => {
+      mocks.search.mockResolvedValue({ documents: excerpts })
+      const pendingB = createDeferred<boolean>()
+      mocks.service.mockResolvedValue({
+        policy: defaultLiveSearchPolicy(),
+        partial: false,
+        verify: async ({ id }: { id: string }) => (id === 'b' ? pendingB.promise : id !== 'c'),
+      })
+      mocks.preview.mockImplementation(
+        async (_provider, _client, { id }: { id: string }) => `Full ${id} launch text`
+      )
+      const previewedIds = () => mocks.preview.mock.calls.map(([, , { id }]) => id)
+      const pending = searchLiveKnowledge.execute({ principal, input })
+      await vi.waitFor(() => expect(previewedIds()).toEqual(['a', 'd']))
+      pendingB.resolve(true)
+      const result = await pending
+      expect(contents(result)).toEqual([
+        'Full a launch text',
+        'Full b launch text',
+        'Full d launch text',
+      ])
+      expect(previewedIds()).toEqual(['a', 'd', 'b'])
+    })
+
+    it('reads no preview until a verification that spends the same request budget finishes', async () => {
+      const googleDoc = 'application/vnd.google-apps.document'
+      vi.mocked(livePolicyFor).mockReturnValue({
+        ...defaultLiveSearchPolicy(),
+        fileTypes: [googleDoc],
+      })
+      mocks.search.mockResolvedValue({ documents: excerpts })
+      const pendingB = createDeferred<unknown>()
+      mocks.json.mockReset().mockImplementation(async (path: string) => {
+        const id = path.split('/').at(-1)
+        return id === 'b' ? pendingB.promise : { id, mimeType: googleDoc }
+      })
+      mocks.preview.mockImplementation(
+        async (_provider, _client, { id }: { id: string }) => `Full ${id} launch text`
+      )
+      const pending = searchLiveKnowledge.execute({ principal, input })
+      await vi.waitFor(() => expect(mocks.json).toHaveBeenCalledTimes(4))
+      expect(mocks.preview).not.toHaveBeenCalled()
+      pendingB.resolve({ id: 'b', mimeType: googleDoc })
+      expect(contents(await pending)).toEqual([
+        'Full a launch text',
+        'Full b launch text',
+        'Full c launch text',
+        'Excerpt d',
+      ])
+    })
+
+    it('never surfaces the preview text of a candidate verification denies', async () => {
+      mocks.search.mockResolvedValue({ documents: excerpts })
+      mocks.service.mockResolvedValue({
+        policy: defaultLiveSearchPolicy(),
+        partial: false,
+        verify: async ({ id }: { id: string }) => id !== 'a',
+      })
+      mocks.preview.mockImplementation(async (_provider, _client, { id }: { id: string }) =>
+        id === 'a' ? 'Private a launch text' : `Full ${id} launch text`
+      )
+      const result = await searchLiveKnowledge.execute({ principal, input })
+      expect(contents(result)).toEqual([
+        'Full b launch text',
+        'Full c launch text',
+        'Full d launch text',
+      ])
+      expect(JSON.stringify(result)).not.toContain('Private a')
+    })
   })
 
   it('gives provider searches the account session cancellation signal', async () => {
@@ -1153,4 +1223,105 @@ describe('authorized live retrieval', () => {
       })
     ).toBe(false)
   })
+
+  const accounts = (count: number) =>
+    Array.from({ length: count }, (_, index) => ({
+      ...account,
+      id: `account-${String(index).padStart(2, '0')}`,
+      displayName: `Drive ${String(index).padStart(2, '0')}`,
+    }))
+
+  it('reports each account beyond the search limit instead of omitting it', async () => {
+    mocks.accounts.mockResolvedValue(accounts(21))
+    const result = await searchLiveKnowledge.execute({ principal, input })
+    const statuses = result.live?.accounts ?? []
+    expect(statuses).toHaveLength(21)
+    expect(new Set(statuses.map(({ accountId }) => accountId)).size).toBe(21)
+    expect(statuses.at(-1)).toMatchObject({
+      accountId: 'account-20',
+      status: 'unavailable',
+      message: expect.stringContaining('Not searched'),
+    })
+    expect(result.retrieval.status).toBe('partial')
+  })
+
+  describe('latency', () => {
+    beforeEach(() => {
+      vi.useFakeTimers()
+      vi.spyOn(AbortSignal, 'timeout').mockImplementation((milliseconds) => {
+        const controller = new AbortController()
+        setTimeout(
+          () => controller.abort(new DOMException('deadline', 'TimeoutError')),
+          milliseconds
+        )
+        return controller.signal
+      })
+    })
+    afterEach(() => {
+      vi.mocked(AbortSignal.timeout).mockRestore()
+      vi.useRealTimers()
+    })
+
+    /** Settles after `ms` of fake time, like a provider round trip. */
+    const delayed = <T>(ms: number, value: T) =>
+      new Promise<T>((resolve) => setTimeout(() => resolve(value), ms))
+
+    it('searches every account at once while bounding concurrent account resolution', async () => {
+      mocks.accounts.mockResolvedValue(accounts(20))
+      let resolving = 0
+      let peak = 0
+      mocks.resolveAccount.mockImplementation(async () => {
+        peak = Math.max(peak, ++resolving)
+        await delayed(50, undefined)
+        resolving--
+        return { account, accessToken: 'secret' }
+      })
+      mocks.search.mockImplementation(() => delayed(3000, { documents: [document] }))
+      let result: Awaited<ReturnType<typeof searchLiveKnowledge.execute>> | undefined
+      const pending = searchLiveKnowledge.execute({ principal, input }).then((value) => {
+        result = value
+      })
+      await vi.advanceTimersByTimeAsync(3500)
+      expect(result?.live?.accounts.map(({ status }) => status)).toEqual(
+        Array.from({ length: 20 }, () => 'ok')
+      )
+      expect(peak).toBeLessThanOrEqual(6)
+      await pending
+    })
+
+    it('times an account out at its deadline when its session never opens', async () => {
+      mocks.service.mockReturnValue(new Promise(() => {}))
+      let result: Awaited<ReturnType<typeof searchLiveKnowledge.execute>> | undefined
+      void searchLiveKnowledge.execute({ principal, input }).then((value) => {
+        result = value
+      })
+      await vi.advanceTimersByTimeAsync(12_000)
+      expect(result?.live?.accounts[0]?.status).toBe('timeout')
+      expect(mocks.search).not.toHaveBeenCalled()
+    })
+
+    it('gives each account its own deadline from its slot, and skips one the search left no time', async () => {
+      mocks.accounts.mockResolvedValue(accounts(13))
+      mocks.resolveAccount.mockImplementation(() =>
+        delayed(9500, { account, accessToken: 'secret' })
+      )
+      mocks.search.mockImplementation(async (_provider, _client, search) => {
+        search.signal?.throwIfAborted()
+        return { documents: [document] }
+      })
+      const pending = searchLiveKnowledge.execute({ principal, input })
+      await vi.advanceTimersByTimeAsync(20_000)
+      const statuses = (await pending).live?.accounts ?? []
+      /** The second wave waits 9.5s for a slot, then still has its full account deadline. */
+      expect(statuses.slice(0, 12).map(({ status }) => status)).toEqual(
+        Array.from({ length: 12 }, () => 'ok')
+      )
+      expect(statuses[12]).toMatchObject({
+        accountId: 'account-12',
+        status: 'timeout',
+        message: expect.stringContaining('Not searched before the search deadline'),
+      })
+      expect(mocks.resolveAccount).toHaveBeenCalledTimes(12)
+    })
+  })
 })