Skip to content

Commit bbca625

Browse files
committed
improvement(search): overlap live search stages and report every unsearched source
1 parent 922b667 commit bbca625

8 files changed

Lines changed: 499 additions & 128 deletions

File tree

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

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -56,6 +56,27 @@ export function mapSettledWithConcurrency<T, R>(
5656
})
5757
}
5858

59+
/**
60+
* A gate that runs at most `limit` of the tasks passed to it at once, starting the rest in call
61+
* order as slots free. Bounds one phase of otherwise concurrent work, such as the part that holds
62+
* database connections, without bounding the rest of it.
63+
*/
64+
export function createConcurrencyLimit(limit: number) {
65+
let active = 0
66+
const waiting: (() => void)[] = []
67+
return async <R>(task: () => Promise<R>): Promise<R> => {
68+
if (active >= limit) await new Promise<void>((resolve) => waiting.push(resolve))
69+
else active++
70+
try {
71+
return await task()
72+
} finally {
73+
const next = waiting.shift()
74+
if (next) next()
75+
else active--
76+
}
77+
}
78+
}
79+
5980
/** Default bound for per-row object-storage materialization fan-out. */
6081
export const MATERIALIZE_CONCURRENCY = 20
6182

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

Lines changed: 8 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@ import {
1111
createNativeClient,
1212
NATIVE_SEARCH_REQUEST_BUDGET,
1313
NativeSearchError,
14+
withRequestSignal,
1415
} from '@/lib/sim-search/live/http'
1516
import { readHubSpotMcp, searchHubSpotMcp } from '@/lib/sim-search/live/hubspot-mcp'
1617
import { readLucidMcp, searchLucidMcp } from '@/lib/sim-search/live/lucid-mcp'
@@ -51,8 +52,11 @@ export interface LiveAccountSession {
5152
/** True only when the document is inside the source boundary and the member may read it. */
5253
verify(document: Reference): Promise<boolean>
5354
read(reference: Reference, filters?: WorkspaceSearchFilters): Promise<NativeDocument>
54-
/** Full text of a verified search result whose provider returned only an excerpt or metadata. */
55-
previewText(document: NativeDocument): Promise<string | undefined>
55+
/**
56+
* Full text of a search result whose provider returned only an excerpt or metadata; `signal`
57+
* cancels the read when its preview is no longer wanted.
58+
*/
59+
previewText(document: NativeDocument, signal: AbortSignal): Promise<string | undefined>
5660
/** Checks a fetched document against the source settings saved now, not at session start. */
5761
verifyCurrent(document: NativeDocument): Promise<boolean>
5862
}
@@ -226,9 +230,9 @@ export async function openLiveAccountSession(
226230
})
227231
return readMcp(reference)
228232
},
229-
async previewText(document) {
233+
async previewText(document, previewSignal) {
230234
if (admin || !client) return undefined
231-
return previewNativeText(provider, client, document)
235+
return previewNativeText(provider, withRequestSignal(client, previewSignal), document)
232236
},
233237
async verifyCurrent(document) {
234238
if (!(await isSearchProviderEnabled(provider, resourceScopeFromOwner(owner)))) return false

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

Lines changed: 134 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
import { ErrorCode, McpError } from '@modelcontextprotocol/sdk/types.js'
2-
import { dbChainMockFns, resetDbChainMock } from '@sim/testing'
2+
import { createDeferred, dbChainMockFns, resetDbChainMock } from '@sim/testing'
33
import { createSessionPrincipal } from '@sim/testing/factories/principal.factory'
44
import { setEnv } from '@sim/testing/mocks/env.mock'
55
import {
@@ -740,6 +740,49 @@ describe('authorized live retrieval', () => {
740740
})
741741
expect(mocks.preview).not.toHaveBeenCalled()
742742
})
743+
744+
it('reads a permitted result’s preview while others verify, and never a denied one’s', async () => {
745+
mocks.search.mockResolvedValue({ documents: excerpts })
746+
const pendingB = createDeferred<boolean>()
747+
mocks.service.mockResolvedValue({
748+
policy: defaultLiveSearchPolicy(),
749+
partial: false,
750+
verify: async ({ id }: { id: string }) => (id === 'b' ? pendingB.promise : id !== 'c'),
751+
})
752+
mocks.preview.mockImplementation(
753+
async (_provider, _client, { id }: { id: string }) => `Full ${id} launch text`
754+
)
755+
const previewedIds = () => mocks.preview.mock.calls.map(([, , { id }]) => id)
756+
const pending = searchLiveKnowledge.execute({ principal, input })
757+
await vi.waitFor(() => expect(previewedIds()).toEqual(['a', 'd']))
758+
pendingB.resolve(true)
759+
const result = await pending
760+
expect(contents(result)).toEqual([
761+
'Full a launch text',
762+
'Full b launch text',
763+
'Full d launch text',
764+
])
765+
expect(previewedIds()).toEqual(['a', 'd', 'b'])
766+
})
767+
768+
it('never surfaces the preview text of a candidate verification denies', async () => {
769+
mocks.search.mockResolvedValue({ documents: excerpts })
770+
mocks.service.mockResolvedValue({
771+
policy: defaultLiveSearchPolicy(),
772+
partial: false,
773+
verify: async ({ id }: { id: string }) => id !== 'a',
774+
})
775+
mocks.preview.mockImplementation(async (_provider, _client, { id }: { id: string }) =>
776+
id === 'a' ? 'Private a launch text' : `Full ${id} launch text`
777+
)
778+
const result = await searchLiveKnowledge.execute({ principal, input })
779+
expect(contents(result)).toEqual([
780+
'Full b launch text',
781+
'Full c launch text',
782+
'Full d launch text',
783+
])
784+
expect(JSON.stringify(result)).not.toContain('Private a')
785+
})
743786
})
744787

745788
it('gives provider searches the account session cancellation signal', async () => {
@@ -1153,4 +1196,94 @@ describe('authorized live retrieval', () => {
11531196
})
11541197
).toBe(false)
11551198
})
1199+
1200+
const accounts = (count: number) =>
1201+
Array.from({ length: count }, (_, index) => ({
1202+
...account,
1203+
id: `account-${String(index).padStart(2, '0')}`,
1204+
displayName: `Drive ${String(index).padStart(2, '0')}`,
1205+
}))
1206+
1207+
it('reports each account beyond the search limit instead of omitting it', async () => {
1208+
mocks.accounts.mockResolvedValue(accounts(21))
1209+
const result = await searchLiveKnowledge.execute({ principal, input })
1210+
const statuses = result.live?.accounts ?? []
1211+
expect(statuses).toHaveLength(21)
1212+
expect(new Set(statuses.map(({ accountId }) => accountId)).size).toBe(21)
1213+
expect(statuses.at(-1)).toMatchObject({
1214+
accountId: 'account-20',
1215+
status: 'unavailable',
1216+
message: expect.stringContaining('Not searched'),
1217+
})
1218+
expect(result.retrieval.status).toBe('partial')
1219+
})
1220+
1221+
describe('latency', () => {
1222+
beforeEach(() => {
1223+
vi.useFakeTimers()
1224+
vi.spyOn(AbortSignal, 'timeout').mockImplementation((milliseconds) => {
1225+
const controller = new AbortController()
1226+
setTimeout(
1227+
() => controller.abort(new DOMException('deadline', 'TimeoutError')),
1228+
milliseconds
1229+
)
1230+
return controller.signal
1231+
})
1232+
})
1233+
afterEach(() => {
1234+
vi.mocked(AbortSignal.timeout).mockRestore()
1235+
vi.useRealTimers()
1236+
})
1237+
1238+
/** Settles after `ms` of fake time, like a provider round trip. */
1239+
const delayed = <T>(ms: number, value: T) =>
1240+
new Promise<T>((resolve) => setTimeout(() => resolve(value), ms))
1241+
1242+
it('searches every account at once while bounding concurrent account resolution', async () => {
1243+
mocks.accounts.mockResolvedValue(accounts(20))
1244+
let resolving = 0
1245+
let peak = 0
1246+
mocks.resolveAccount.mockImplementation(async () => {
1247+
peak = Math.max(peak, ++resolving)
1248+
await delayed(50, undefined)
1249+
resolving--
1250+
return { account, accessToken: 'secret' }
1251+
})
1252+
mocks.search.mockImplementation(() => delayed(3000, { documents: [document] }))
1253+
let result: Awaited<ReturnType<typeof searchLiveKnowledge.execute>> | undefined
1254+
const pending = searchLiveKnowledge.execute({ principal, input }).then((value) => {
1255+
result = value
1256+
})
1257+
await vi.advanceTimersByTimeAsync(3500)
1258+
expect(result?.live?.accounts.map(({ status }) => status)).toEqual(
1259+
Array.from({ length: 20 }, () => 'ok')
1260+
)
1261+
expect(peak).toBeLessThanOrEqual(6)
1262+
await pending
1263+
})
1264+
1265+
it('gives each account its own deadline from its slot, and skips one the search left no time', async () => {
1266+
mocks.accounts.mockResolvedValue(accounts(13))
1267+
mocks.resolveAccount.mockImplementation(() =>
1268+
delayed(9500, { account, accessToken: 'secret' })
1269+
)
1270+
mocks.search.mockImplementation(async (_provider, _client, search) => {
1271+
search.signal?.throwIfAborted()
1272+
return { documents: [document] }
1273+
})
1274+
const pending = searchLiveKnowledge.execute({ principal, input })
1275+
await vi.advanceTimersByTimeAsync(20_000)
1276+
const statuses = (await pending).live?.accounts ?? []
1277+
/** The second wave waits 9.5s for a slot, then still has its full account deadline. */
1278+
expect(statuses.slice(0, 12).map(({ status }) => status)).toEqual(
1279+
Array.from({ length: 12 }, () => 'ok')
1280+
)
1281+
expect(statuses[12]).toMatchObject({
1282+
accountId: 'account-12',
1283+
status: 'timeout',
1284+
message: expect.stringContaining('Not searched before the search deadline'),
1285+
})
1286+
expect(mocks.resolveAccount).toHaveBeenCalledTimes(12)
1287+
})
1288+
})
11561289
})

0 commit comments

Comments
 (0)