diff --git a/apps/sim/lib/core/utils/concurrency.ts b/apps/sim/lib/core/utils/concurrency.ts index 70b9e701522..9e69fd4e5ff 100644 --- a/apps/sim/lib/core/utils/concurrency.ts +++ b/apps/sim/lib/core/utils/concurrency.ts @@ -56,6 +56,27 @@ export function mapSettledWithConcurrency( }) } +/** + * 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 (task: () => Promise): Promise => { + if (active >= limit) await new Promise((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 diff --git a/apps/sim/lib/sim-search/live/account-session.ts b/apps/sim/lib/sim-search/live/account-session.ts index b8fede59474..ba3c256e6a0 100644 --- a/apps/sim/lib/sim-search/live/account-session.ts +++ b/apps/sim/lib/sim-search/live/account-session.ts @@ -11,6 +11,7 @@ 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' @@ -18,7 +19,7 @@ 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 read(reference: Reference, filters?: WorkspaceSearchFilters): Promise - /** Full text of a verified search result whose provider returned only an excerpt or metadata. */ - previewText(document: NativeDocument): Promise + /** + * 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 + /** + * 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 } @@ -174,7 +183,7 @@ 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, }) @@ -182,6 +191,7 @@ export async function openLiveAccountSession( 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 => { 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 diff --git a/apps/sim/lib/sim-search/live/application.test.ts b/apps/sim/lib/sim-search/live/application.test.ts index ff576950b77..466ccd94d4c 100644 --- a/apps/sim/lib/sim-search/live/application.test.ts +++ b/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() + 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() + 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 = (ms: number, value: T) => + new Promise((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> | 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> | 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) + }) + }) }) diff --git a/apps/sim/lib/sim-search/live/application.ts b/apps/sim/lib/sim-search/live/application.ts index ad446352c0f..07595b70402 100644 --- a/apps/sim/lib/sim-search/live/application.ts +++ b/apps/sim/lib/sim-search/live/application.ts @@ -27,7 +27,11 @@ import { resourceScopeKey, } from '@/lib/core/resource-scope' import { createPinnedConnectionPool } from '@/lib/core/security/input-validation.server' -import { mapWithConcurrency, waitWithAbort } from '@/lib/core/utils/concurrency' +import { + createConcurrencyLimit, + mapWithConcurrency, + waitWithAbort, +} from '@/lib/core/utils/concurrency' import { MANAGED_MCP_CONNECTORS } from '@/lib/credential-groups/managed-mcp-connectors' import { requireOrganizationSearchAvailable } from '@/lib/knowledge/access/availability' import { defineAuthorizedKnowledgeUseCase } from '@/lib/knowledge/application/authorized-knowledge-use-case' @@ -279,11 +283,19 @@ function resultFor( } /** - * Accounts searched at once; bounds token refresh and provider fan-out in large organizations. - * Each account is one provider, and a search waits for every account, so a second wave adds a - * full provider round trip; eight covers the account count most organizations search. + * Accounts resolving their credential and opening their session at once. That phase reads the + * database (credentials, source settings, provider rollout) along with the network calls those + * reads lead to, such as a token refresh or a Google Directory lookup and delegated token mint, + * so it is bounded well under the web database pool. Provider searches are not bounded: a search + * waits for every account, and a second wave would add a full provider round trip. */ -const ACCOUNT_CONCURRENCY = 8 +const ACCOUNT_OPEN_CONCURRENCY = 6 +/** Milliseconds a whole search may take. */ +const SEARCH_DEADLINE_MS = 20_000 +/** Milliseconds one account may take once it gets a slot, within what remains of the search. */ +const ACCOUNT_DEADLINE_MS = 12_000 +/** An account reaching its turn with less than this left before its deadline is not searched. */ +const ACCOUNT_START_FLOOR_MS = 1500 /** Candidates verified at once; each verification is one or more provider requests. */ const VERIFY_CONCURRENCY = 5 const MAX_ACCOUNTS = 20 @@ -296,21 +308,24 @@ const READ_WINDOW_CHARACTERS = 8000 * Verifies candidates in rank order. A revoked grant fails the whole account; any other * verification failure only omits that candidate, and is reported as incomplete coverage. */ -async function verifyCandidates(session: LiveAccountSession, candidates: LiveCandidate[]) { - const outcomes = await mapWithConcurrency( - candidates, - VERIFY_CONCURRENCY, - async ({ document }) => { - try { - return (await session.verify(document)) ? ('permitted' as const) : ('denied' as const) - } catch (error) { +async function verifyCandidates( + session: LiveAccountSession, + candidates: LiveCandidate[], + onVerified?: (candidate: LiveCandidate, permitted: boolean) => void +) { + const outcomes = await mapWithConcurrency(candidates, VERIFY_CONCURRENCY, async (candidate) => { + const outcome = await session.verify(candidate.document).then( + (permitted) => (permitted ? ('permitted' as const) : ('denied' as const)), + (error: unknown) => { if (error instanceof NativeSearchError && error.status === 'reconnect') throw error return error instanceof NativeSearchError && error.status === 'rate_limited' ? ('rate_limited' as const) : ('unverified' as const) } - } - ) + ) + onVerified?.(candidate, outcome === 'permitted') + return outcome + }) return { permitted: candidates.filter((_, index) => outcomes[index] === 'permitted'), unverified: outcomes.some((outcome) => outcome === 'unverified' || outcome === 'rate_limited'), @@ -328,30 +343,55 @@ const PREVIEW_SOURCE_CHARACTERS = 100_000 /** * Replaces the excerpt or metadata of an account's top results with their full text, so their * preview is a passage around the match like other providers' results rather than a prompt to - * read. A read that fails or misses the deadline keeps the provider's own preview. + * read. A result is read only once verification permits it and it is among the top results the + * denials so far leave, so reads overlap the verification of other results without reading one + * that is denied. A session whose verification spends the budget reads draw on reads only after + * verification, so a preview never costs a result its check. `attach` gives text only to the verified results it is passed and cancels every + * other read. A read that fails or misses its deadline keeps the provider's preview. */ -async function withPreviewText( +function streamPreviewReads( session: LiveAccountSession, candidates: LiveCandidate[], + matchingOf: (verified: LiveCandidate[]) => LiveCandidate[], query: string -): Promise { - const deadline = AbortSignal.timeout(PREVIEW_DEADLINE_MS) - const expired = new Promise((resolve) => - deadline.addEventListener('abort', () => resolve(undefined), { once: true }) - ) - const texts = await Promise.all( - candidates - .slice(0, PREVIEW_READS) - .map((candidate) => - Promise.race([session.previewText(candidate.document).catch(() => undefined), expired]) - ) - ) - return candidates.map((candidate, index) => { - const text = texts[index]?.trim() - if (!text) return candidate - const content = matchPassage(text, query, PREVIEW_SOURCE_CHARACTERS).content - return { ...candidate, document: { ...candidate.document, content } } - }) +) { + const reads = new Map; abandoned: AbortController }>() + const permitted = new Set() + const rejected = new Set() + const read = (candidate: LiveCandidate) => { + const started = reads.get(candidate.documentId) + if (started) return started.text + const abandoned = new AbortController() + const signal = AbortSignal.any([abandoned.signal, AbortSignal.timeout(PREVIEW_DEADLINE_MS)]) + const text = waitWithAbort(session.previewText(candidate.document, signal), signal).catch( + () => undefined + ) + reads.set(candidate.documentId, { text, abandoned }) + return text + } + return { + verified(candidate: LiveCandidate, allowed: boolean) { + ;(allowed ? permitted : rejected).add(candidate.documentId) + const remaining = candidates.filter(({ documentId }) => !rejected.has(documentId)) + for (const leading of matchingOf(remaining).slice(0, PREVIEW_READS)) + if (permitted.has(leading.documentId)) read(leading) + }, + async attach(verified: LiveCandidate[]): Promise { + const leading = verified.slice(0, PREVIEW_READS) + const kept = new Set(leading.map(({ documentId }) => documentId)) + for (const [documentId, { abandoned }] of reads) if (!kept.has(documentId)) abandoned.abort() + const texts = await Promise.all(leading.map(read)) + return verified.map((candidate, index) => { + const text = texts[index]?.trim() + if (!text) return candidate + const content = matchPassage(text, query, PREVIEW_SOURCE_CHARACTERS).content + return { ...candidate, document: { ...candidate.document, content } } + }) + }, + cancel() { + for (const { abandoned } of reads.values()) abandoned.abort() + }, + } } /** Keeps the first of candidates a provider marked as the same item, such as a shared meeting. */ @@ -439,9 +479,10 @@ const liveKnowledgeSearch = defineAuthorizedKnowledgeUseCase({ throw new OrchestrationError('validation', 'endDate must be after startDate') if (queries?.some((query) => !query.query && !query.browse) && !hasDateBounds(filters)) throw new OrchestrationError('validation', 'Empty native queries require a date bound') + const searchDeadline = Date.now() + SEARCH_DEADLINE_MS const searchSignal = input.signal - ? AbortSignal.any([input.signal, AbortSignal.timeout(20_000)]) - : AbortSignal.timeout(20_000) + ? AbortSignal.any([input.signal, AbortSignal.timeout(SEARCH_DEADLINE_MS)]) + : AbortSignal.timeout(SEARCH_DEADLINE_MS) /** An account's native queries with their request index; none searches it with the plain query. */ const nativesFor = (account: LiveAccount): NativeTarget[] => queries @@ -461,6 +502,7 @@ const liveKnowledgeSearch = defineAuthorizedKnowledgeUseCase({ ) .sort(compareAccounts) const selected = eligible.slice(0, MAX_ACCOUNTS) + const openingAccount = createConcurrencyLimit(ACCOUNT_OPEN_CONCURRENCY) annotateSearchDiagnostics({ operation: 'search_workspace', accountCount: eligible.length, @@ -543,32 +585,45 @@ const liveKnowledgeSearch = defineAuthorizedKnowledgeUseCase({ const candidates = page.documents .filter((document) => document.id) .map((document) => candidateFor(document, account, input, userId)) + const matchingOf = (verified: LiveCandidate[]) => + firstOfEachDocument( + verified.filter(({ document, documentId }) => + matchesLiveFilters(document, documentId, account.provider, queryFilters) + ) + ) /** * Local filters run first so provider verification is spent only on eligible results. * Undated results are verified too, so their exclusion is reported only when readable. */ - const { permitted, unverified, rateLimited } = await measureSearchStage('live.verify', () => - verifyCandidates( - session, - candidates.filter( - ({ document, documentId }) => - matchesLiveFilters(document, documentId, account.provider, queryFilters) || - lacksFilterDate(document, account.provider, queryFilters) - ) - ) - ) - const matching = firstOfEachDocument( - permitted.filter(({ document, documentId }) => - matchesLiveFilters(document, documentId, account.provider, queryFilters) - ) + const eligible = candidates.filter( + ({ document, documentId }) => + matchesLiveFilters(document, documentId, account.provider, queryFilters) || + lacksFilterDate(document, account.provider, queryFilters) ) /** A native query's operators would anchor the passage, so only plain queries preview. */ - const previewed = + const previews = !native && input.query.trim() - ? await measureSearchStage('live.preview', () => - withPreviewText(session, matching, input.query) - ) + ? streamPreviewReads(session, eligible, matchingOf, input.query) + : undefined + let verification: Awaited> + let matching: LiveCandidate[] + let previewed: LiveCandidate[] + try { + verification = await measureSearchStage('live.verify', () => + verifyCandidates( + session, + eligible, + session.verifySharesPreviewBudget ? undefined : previews?.verified + ) + ) + matching = matchingOf(verification.permitted) + previewed = previews + ? await measureSearchStage('live.preview', () => previews.attach(matching)) : matching + } finally { + previews?.cancel() + } + const { permitted, unverified, rateLimited } = verification const undatedExcluded = permitted.some(({ document }) => lacksFilterDate(document, account.provider, queryFilters) ) @@ -684,7 +739,9 @@ const liveKnowledgeSearch = defineAuthorizedKnowledgeUseCase({ }) /** Cancels requests still in flight once the account settles, including after a failure. */ const settled = new AbortController() - const signal = AbortSignal.any([searchSignal, AbortSignal.timeout(12_000), settled.signal]) + /** Aborts at the account's own deadline, which starts once it gets a slot to open. */ + const deadline = new AbortController() + const signal = AbortSignal.any([searchSignal, deadline.signal, settled.signal]) const failed = (error: unknown, target: NativeTarget): SearchedQuery => { input.signal?.throwIfAborted() const failure = @@ -715,13 +772,21 @@ const liveKnowledgeSearch = defineAuthorizedKnowledgeUseCase({ } let session: LiveAccountSession | undefined try { - signal.throwIfAborted() - /** Bounds every resolution path, including those that take no signal (admin GitLab, managed OAuth). */ - const resolved = await measureSearchStage('live.resolve', () => - waitWithAbort(resolveListedLiveAccount(input, userId, account, signal), signal) - ) - session = await measureSearchStage('live.session', () => - openLiveAccountSession({ + const opened = await openingAccount(async () => { + const now = Date.now() + const remaining = Math.min(searchDeadline, now + ACCOUNT_DEADLINE_MS) - now + if (remaining < ACCOUNT_START_FLOOR_MS) return undefined + const expiry = AbortSignal.timeout(remaining) + expiry.addEventListener('abort', () => deadline.abort(expiry.reason), { + once: true, + signal: settled.signal, + }) + signal.throwIfAborted() + /** Bounds every resolution path, including those that take no signal (admin GitLab, managed OAuth). */ + const resolved = await measureSearchStage('live.resolve', () => + waitWithAbort(resolveListedLiveAccount(input, userId, account, signal), signal) + ) + const opening = openLiveAccountSession({ owner: input, userId, resolved, @@ -730,8 +795,26 @@ const liveKnowledgeSearch = defineAuthorizedKnowledgeUseCase({ pool, searches: natives.length, }) - ) - const currentSession = session + /** Opening reads the database, which takes no signal; a session that opens late is closed. */ + session = await measureSearchStage('live.session', () => + waitWithAbort(opening, signal) + ).catch((error: unknown) => { + opening.then((late) => late.close()).catch(() => {}) + throw error + }) + return { resolved, session } + }) + if (!opened) + return natives.map((target) => + failed( + new NativeSearchError( + 'timeout', + 'Not searched before the search deadline. Search this account alone or narrow the query.' + ), + target + ) + ) + const { resolved, session: currentSession } = opened return await Promise.all( natives.map((target) => withSearchQueryDiagnostics(target.queryIndex, () => @@ -749,9 +832,11 @@ const liveKnowledgeSearch = defineAuthorizedKnowledgeUseCase({ let searched: SearchedQuery[] try { searched = ( - await mapWithConcurrency(selected, ACCOUNT_CONCURRENCY, (account, accountIndex) => - withSearchProviderDiagnostics({ provider: account.provider, accountIndex }, () => - searchAccount(account) + await Promise.all( + selected.map((account, accountIndex) => + withSearchProviderDiagnostics({ provider: account.provider, accountIndex }, () => + searchAccount(account) + ) ) ) ).flat() @@ -818,8 +903,23 @@ const liveKnowledgeSearch = defineAuthorizedKnowledgeUseCase({ ]), } }) + /** Live coverage guides the providers this call searched or tried to search. */ + const guidance = liveSearchGuidance( + accounts.filter((account) => account.status !== 'reconnect').map(({ provider }) => provider) + ) + for (const account of eligible.slice(MAX_ACCOUNTS)) { + for (const { queryIndex } of nativesFor(account)) + accounts.push({ + accountId: account.id, + provider: account.provider, + displayName: account.displayName, + ...(queryIndex === undefined ? {} : { queryIndex }), + status: 'unavailable', + message: `Not searched: one search covers at most ${MAX_ACCOUNTS} accounts. Search this source or account alone.`, + }) + } for (const [queryIndex, query] of (queries ?? []).entries()) { - if (!selected.some((account) => targetsAccount(query, account))) + if (!eligible.some((account) => targetsAccount(query, account))) accounts.push({ accountId: query.accountId ?? '', provider: query.provider, @@ -829,10 +929,9 @@ const liveKnowledgeSearch = defineAuthorizedKnowledgeUseCase({ message: 'No connection with this provider is configured and approved in this scope.', }) } - const retrievalStatus = - accounts.some((account) => account.status !== 'ok') || eligible.length > selected.length - ? 'partial' - : 'complete' + const retrievalStatus = accounts.some((account) => account.status !== 'ok') + ? 'partial' + : 'complete' annotateSearchDiagnostics({ resultCount: Math.min(ranked.length, input.topK), retrievalStatus, @@ -847,11 +946,7 @@ const liveKnowledgeSearch = defineAuthorizedKnowledgeUseCase({ live: { backend: 'live', accounts, - guidance: liveSearchGuidance( - accounts - .filter((account) => account.status !== 'reconnect') - .map(({ provider }) => provider) - ), + guidance, }, } }, diff --git a/apps/sim/lib/sim-search/live/github-service.ts b/apps/sim/lib/sim-search/live/github-service.ts index 7d2f2645e7d..9a2021bc9f3 100644 --- a/apps/sim/lib/sim-search/live/github-service.ts +++ b/apps/sim/lib/sim-search/live/github-service.ts @@ -132,6 +132,7 @@ export function createGitHubServiceVerifier( return { policy, partial: false, + spendsMemberBudget: true, async verify(document: Pick) { const repository = document.container?.toLowerCase() if (!repository) return false diff --git a/apps/sim/lib/sim-search/live/google-service.test.ts b/apps/sim/lib/sim-search/live/google-service.test.ts index b18b6336f88..545faed5a67 100644 --- a/apps/sim/lib/sim-search/live/google-service.test.ts +++ b/apps/sim/lib/sim-search/live/google-service.test.ts @@ -232,3 +232,93 @@ describe('Google service search scope and continuation', () => { }) }) }) + +describe('Google service session latency', () => { + /** Settles after `ms` of fake time, like a provider round trip. */ + const delayed = (ms: number, value: T) => + new Promise((resolve) => setTimeout(() => resolve(value), ms)) + const failAfter = (ms: number, error: Error) => + new Promise((_, reject) => setTimeout(() => reject(error), ms)) + const clientFor = (json: NativeClient['json']) => ({ + json: vi.fn(json), + text: vi.fn(), + bytes: vi.fn(), + }) + + it('looks up the member identity, the administrator, and every source user concurrently', async () => { + member.json.mockImplementation(() => + delayed(1000, { user: { emailAddress: 'reader@example.com' } }) + ) + mocks.directory.mockImplementation((_token, email: string) => delayed(1000, person(email))) + mint.mockImplementation(() => delayed(1000, 'delegated-token')) + let settled = false + const session = create('google_drive', { + userEmails: ['first@example.com', 'second@example.com', 'third@example.com'], + }).finally(() => { + settled = true + }) + /** Identity and administrator in parallel, then each user's lookup and mint in parallel. */ + await vi.advanceTimersByTimeAsync(3000) + expect(settled).toBe(true) + await vi.runAllTimersAsync() + expect((await session).partial).toBe(false) + expect(mint.mock.calls.map(([email]) => email)).toEqual([ + 'first@example.com', + 'second@example.com', + 'third@example.com', + ]) + }) + + it('surfaces the member identity failure ahead of a failed administrator lookup', async () => { + member.json.mockImplementation(() => + failAfter(500, new NativeSearchError('reconnect', 'Member grant revoked')) + ) + mocks.directory.mockRejectedValue(new Error('Directory unavailable')) + const session = create('gmail') + const outcome = expect(session).rejects.toThrow('Member grant revoked') + await vi.advanceTimersByTimeAsync(500) + await outcome + }) + + it('verifies a Drive file with one source user at a time, stopping at the first allow', async () => { + member.json.mockResolvedValue({ user: { emailAddress: 'reader@example.com' } }) + const first = clientFor(async () => ({ id: 'file', mimeType: 'text/plain' })) + const second = clientFor(async () => ({ id: 'file', mimeType: 'text/plain' })) + mocks.createClient.mockReturnValueOnce(first).mockReturnValueOnce(second) + const session = await create('google_drive', { + userEmails: ['first@example.com', 'second@example.com'], + }) + expect(await session.verify({ id: 'file' })).toBe(true) + expect(first.json).toHaveBeenCalledOnce() + expect(second.json).not.toHaveBeenCalled() + }) + + it('fails on a throttled source user before asking a later one', async () => { + member.json.mockResolvedValue({ user: { emailAddress: 'reader@example.com' } }) + const throttled = clientFor(() => + Promise.reject(new NativeSearchError('rate_limited', 'Provider limit')) + ) + const allowed = clientFor(async () => ({ id: 'file', mimeType: 'text/plain' })) + mocks.createClient.mockReturnValueOnce(throttled).mockReturnValueOnce(allowed) + const session = await create('google_drive', { + userEmails: ['first@example.com', 'second@example.com'], + }) + await expect(session.verify({ id: 'file' })).rejects.toThrow('Provider limit') + expect(allowed.json).not.toHaveBeenCalled() + }) + + it('propagates throttling over an unavailable user when no source user allows the file', async () => { + member.json.mockResolvedValue({ user: { emailAddress: 'reader@example.com' } }) + mocks.createClient + .mockReturnValueOnce( + clientFor(() => Promise.reject(new NativeSearchError('unavailable', 'Source timed out'))) + ) + .mockReturnValueOnce( + clientFor(() => Promise.reject(new NativeSearchError('rate_limited', 'Provider limit'))) + ) + const session = await create('google_drive', { + userEmails: ['first@example.com', 'second@example.com'], + }) + await expect(session.verify({ id: 'file' })).rejects.toThrow('Provider limit') + }) +}) diff --git a/apps/sim/lib/sim-search/live/google-service.ts b/apps/sim/lib/sim-search/live/google-service.ts index 67f398d2ace..9843ae06335 100644 --- a/apps/sim/lib/sim-search/live/google-service.ts +++ b/apps/sim/lib/sim-search/live/google-service.ts @@ -24,6 +24,8 @@ const SOURCE_REQUEST_BUDGET = 30 /** Distinct ancestor folders a restricted Drive page may read while verifying its files. */ const DRIVE_FOLDER_RESERVE = 9 const DAY_MS = 86400000 +/** Source users a session delegates to; more are reported as partial coverage. */ +const MAX_SUBJECTS = 6 interface GoogleServiceSession { verify(document: Reference): Promise @@ -49,11 +51,18 @@ export async function createGoogleServiceVerifier(input: { pool?: PinnedConnectionPool }): Promise { const { provider, token, member, config, policy, signal, pool } = input - if (!token.getDelegatedAccessToken) + const { getDelegatedAccessToken } = token + if (!getDelegatedAccessToken) throw new NativeSearchError( 'unavailable', 'Configure a Google Workspace service account with domain-wide delegation.' ) + /** The administrator lookup does not depend on the member identity, so both run at once. */ + const administratorLookup = getGoogleWorkspaceUser(token.accessToken, string(config.adminEmail), { + signal, + }) + // Awaited below; this keeps it handled when the identity lookup fails first. + administratorLookup.catch(() => undefined) const identity = object( await member.json( provider === 'gmail' @@ -78,9 +87,7 @@ export async function createGoogleServiceVerifier(input: { 'reconnect', 'The connected Google account identity could not be verified.' ) - const administrator = await getGoogleWorkspaceUser(token.accessToken, string(config.adminEmail), { - signal, - }) + const administrator = await administratorLookup if (!administrator?.active) throw new NativeSearchError('unavailable', 'The source Directory administrator is unavailable.') const selected = selectedGoogleWorkspaceUsers(config.userEmails) @@ -88,37 +95,42 @@ export async function createGoogleServiceVerifier(input: { if (provider !== 'google_drive' && selected.length && !selected.includes(email)) { return { verify: async () => false, partial: false, scopeSearch: () => null } } - const clients: { - email: string - client: NativeClient - verifyPolicy: ReturnType - }[] = [] - for (const subject of subjects.slice(0, 6)) { - signal.throwIfAborted() - const person = await getGoogleWorkspaceUser(token.accessToken, subject, { signal }) - if ( - !person?.active || - person.customerId !== administrator.customerId || - person.email !== subject - ) - continue - const accessToken = await token.getDelegatedAccessToken(person.email, signal) - const client = createNativeClient({ - origin: 'https://www.googleapis.com', - accessToken, - signal, - pool, - }) - clients.push({ - email: person.email, - client, - verifyPolicy: createPolicyVerifier(provider, policy, client, 'https://www.googleapis.com'), + /** + * Each subject's Directory check and token mint run alongside the others'; results keep the + * configured subject order, and the first failure in that order fails the session. + */ + const delegations = await Promise.allSettled( + subjects.slice(0, MAX_SUBJECTS).map(async (subject) => { + signal.throwIfAborted() + const person = await getGoogleWorkspaceUser(token.accessToken, subject, { signal }) + if ( + !person?.active || + person.customerId !== administrator.customerId || + person.email !== subject + ) + return undefined + const accessToken = await getDelegatedAccessToken(person.email, signal) + const client = createNativeClient({ + origin: 'https://www.googleapis.com', + accessToken, + signal, + pool, + }) + return { + email: person.email, + client, + verifyPolicy: createPolicyVerifier(provider, policy, client, 'https://www.googleapis.com'), + } }) - } - let partial = subjects.length > 6 + ) + const clients = delegations.flatMap((delegation) => { + if (delegation.status === 'rejected') throw delegation.reason + return delegation.value ? [delegation.value] : [] + }) + let partial = subjects.length > MAX_SUBJECTS if (provider === 'google_drive' && !selected.length && !clients.length) { signal.throwIfAborted() - const accessToken = await token.getDelegatedAccessToken(administrator.email, signal) + const accessToken = await getDelegatedAccessToken(administrator.email, signal) const client = createNativeClient({ origin: 'https://www.googleapis.com', accessToken, diff --git a/apps/sim/lib/sim-search/live/http.ts b/apps/sim/lib/sim-search/live/http.ts index f24d814d1f7..eb2d889a444 100644 --- a/apps/sim/lib/sim-search/live/http.ts +++ b/apps/sim/lib/sim-search/live/http.ts @@ -160,9 +160,9 @@ export function createNativeClient(input: { async json(path, options) { return request(path, options, (response) => response.json()) }, - async bytes(path, query) { + async bytes(path, query, options) { try { - return await request(path, { query }, async (response) => + return await request(path, { query, signal: options?.signal }, async (response) => Buffer.from(await response.arrayBuffer()) ) } catch (error) { @@ -170,12 +170,30 @@ export function createNativeClient(input: { throw error } }, - async text(path, query) { - return request(path, { query }, (response) => response.text()) + async text(path, query, options) { + return request(path, { query, signal: options?.signal }, (response) => response.text()) }, }) } +/** + * The client with `signal` cancelling each of its requests, for work that may be abandoned + * before the client's own operation ends. A memoized GET keeps only the client's signal, since + * its response is shared with callers that did not abandon it. + */ +export function withRequestSignal(client: NativeClient, signal: AbortSignal): NativeClient { + const scoped = (own?: AbortSignal) => (own ? AbortSignal.any([own, signal]) : signal) + return { + json: (path, options) => + client.json( + path, + options?.memo && !options.body ? options : { ...options, signal: scoped(options?.signal) } + ), + text: (path, query, options) => client.text(path, query, { signal: scoped(options?.signal) }), + bytes: (path, query, options) => client.bytes(path, query, { signal: scoped(options?.signal) }), + } +} + /** * Answers a repeated `memo` GET from the client's earlier response, including a failure, so a * rate-limited endpoint is not requested again. A client lives for one search or read, so a diff --git a/apps/sim/lib/sim-search/live/service-session.ts b/apps/sim/lib/sim-search/live/service-session.ts index 30c8601d79d..df68caa9a5a 100644 --- a/apps/sim/lib/sim-search/live/service-session.ts +++ b/apps/sim/lib/sim-search/live/service-session.ts @@ -30,6 +30,8 @@ import { CONNECTOR_META_REGISTRY } from '@/connectors/registry' interface LiveServiceSession { policy: LiveSearchPolicy partial: boolean + /** True when verification also spends requests from the member client's budget. */ + spendsMemberBudget?: boolean verify(document: Pick): Promise scopeSearch?(input: NativeSearchInput): NativeSearchInput | null } diff --git a/apps/sim/lib/sim-search/live/types.ts b/apps/sim/lib/sim-search/live/types.ts index a9cb8baee12..a573634cf97 100644 --- a/apps/sim/lib/sim-search/live/types.ts +++ b/apps/sim/lib/sim-search/live/types.ts @@ -86,8 +86,16 @@ export interface NativeClient { signal?: AbortSignal } ): Promise - text(path: string, query?: Record): Promise - bytes(path: string, query?: Record): Promise + text( + path: string, + query?: Record, + options?: { signal?: AbortSignal } + ): Promise + bytes( + path: string, + query?: Record, + options?: { signal?: AbortSignal } + ): Promise } export interface NativeSearchInput {