|
1 | 1 | import { ErrorCode, McpError } from '@modelcontextprotocol/sdk/types.js' |
2 | | -import { dbChainMockFns, resetDbChainMock } from '@sim/testing' |
| 2 | +import { createDeferred, dbChainMockFns, resetDbChainMock } from '@sim/testing' |
3 | 3 | import { createSessionPrincipal } from '@sim/testing/factories/principal.factory' |
4 | 4 | import { setEnv } from '@sim/testing/mocks/env.mock' |
5 | 5 | import { |
@@ -740,6 +740,76 @@ describe('authorized live retrieval', () => { |
740 | 740 | }) |
741 | 741 | expect(mocks.preview).not.toHaveBeenCalled() |
742 | 742 | }) |
| 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('reads no preview until a verification that spends the same request budget finishes', async () => { |
| 769 | + const googleDoc = 'application/vnd.google-apps.document' |
| 770 | + vi.mocked(livePolicyFor).mockReturnValue({ |
| 771 | + ...defaultLiveSearchPolicy(), |
| 772 | + fileTypes: [googleDoc], |
| 773 | + }) |
| 774 | + mocks.search.mockResolvedValue({ documents: excerpts }) |
| 775 | + const pendingB = createDeferred<unknown>() |
| 776 | + mocks.json.mockReset().mockImplementation(async (path: string) => { |
| 777 | + const id = path.split('/').at(-1) |
| 778 | + return id === 'b' ? pendingB.promise : { id, mimeType: googleDoc } |
| 779 | + }) |
| 780 | + mocks.preview.mockImplementation( |
| 781 | + async (_provider, _client, { id }: { id: string }) => `Full ${id} launch text` |
| 782 | + ) |
| 783 | + const pending = searchLiveKnowledge.execute({ principal, input }) |
| 784 | + await vi.waitFor(() => expect(mocks.json).toHaveBeenCalledTimes(4)) |
| 785 | + expect(mocks.preview).not.toHaveBeenCalled() |
| 786 | + pendingB.resolve({ id: 'b', mimeType: googleDoc }) |
| 787 | + expect(contents(await pending)).toEqual([ |
| 788 | + 'Full a launch text', |
| 789 | + 'Full b launch text', |
| 790 | + 'Full c launch text', |
| 791 | + 'Excerpt d', |
| 792 | + ]) |
| 793 | + }) |
| 794 | + |
| 795 | + it('never surfaces the preview text of a candidate verification denies', async () => { |
| 796 | + mocks.search.mockResolvedValue({ documents: excerpts }) |
| 797 | + mocks.service.mockResolvedValue({ |
| 798 | + policy: defaultLiveSearchPolicy(), |
| 799 | + partial: false, |
| 800 | + verify: async ({ id }: { id: string }) => id !== 'a', |
| 801 | + }) |
| 802 | + mocks.preview.mockImplementation(async (_provider, _client, { id }: { id: string }) => |
| 803 | + id === 'a' ? 'Private a launch text' : `Full ${id} launch text` |
| 804 | + ) |
| 805 | + const result = await searchLiveKnowledge.execute({ principal, input }) |
| 806 | + expect(contents(result)).toEqual([ |
| 807 | + 'Full b launch text', |
| 808 | + 'Full c launch text', |
| 809 | + 'Full d launch text', |
| 810 | + ]) |
| 811 | + expect(JSON.stringify(result)).not.toContain('Private a') |
| 812 | + }) |
743 | 813 | }) |
744 | 814 |
|
745 | 815 | it('gives provider searches the account session cancellation signal', async () => { |
@@ -1153,4 +1223,105 @@ describe('authorized live retrieval', () => { |
1153 | 1223 | }) |
1154 | 1224 | ).toBe(false) |
1155 | 1225 | }) |
| 1226 | + |
| 1227 | + const accounts = (count: number) => |
| 1228 | + Array.from({ length: count }, (_, index) => ({ |
| 1229 | + ...account, |
| 1230 | + id: `account-${String(index).padStart(2, '0')}`, |
| 1231 | + displayName: `Drive ${String(index).padStart(2, '0')}`, |
| 1232 | + })) |
| 1233 | + |
| 1234 | + it('reports each account beyond the search limit instead of omitting it', async () => { |
| 1235 | + mocks.accounts.mockResolvedValue(accounts(21)) |
| 1236 | + const result = await searchLiveKnowledge.execute({ principal, input }) |
| 1237 | + const statuses = result.live?.accounts ?? [] |
| 1238 | + expect(statuses).toHaveLength(21) |
| 1239 | + expect(new Set(statuses.map(({ accountId }) => accountId)).size).toBe(21) |
| 1240 | + expect(statuses.at(-1)).toMatchObject({ |
| 1241 | + accountId: 'account-20', |
| 1242 | + status: 'unavailable', |
| 1243 | + message: expect.stringContaining('Not searched'), |
| 1244 | + }) |
| 1245 | + expect(result.retrieval.status).toBe('partial') |
| 1246 | + }) |
| 1247 | + |
| 1248 | + describe('latency', () => { |
| 1249 | + beforeEach(() => { |
| 1250 | + vi.useFakeTimers() |
| 1251 | + vi.spyOn(AbortSignal, 'timeout').mockImplementation((milliseconds) => { |
| 1252 | + const controller = new AbortController() |
| 1253 | + setTimeout( |
| 1254 | + () => controller.abort(new DOMException('deadline', 'TimeoutError')), |
| 1255 | + milliseconds |
| 1256 | + ) |
| 1257 | + return controller.signal |
| 1258 | + }) |
| 1259 | + }) |
| 1260 | + afterEach(() => { |
| 1261 | + vi.mocked(AbortSignal.timeout).mockRestore() |
| 1262 | + vi.useRealTimers() |
| 1263 | + }) |
| 1264 | + |
| 1265 | + /** Settles after `ms` of fake time, like a provider round trip. */ |
| 1266 | + const delayed = <T>(ms: number, value: T) => |
| 1267 | + new Promise<T>((resolve) => setTimeout(() => resolve(value), ms)) |
| 1268 | + |
| 1269 | + it('searches every account at once while bounding concurrent account resolution', async () => { |
| 1270 | + mocks.accounts.mockResolvedValue(accounts(20)) |
| 1271 | + let resolving = 0 |
| 1272 | + let peak = 0 |
| 1273 | + mocks.resolveAccount.mockImplementation(async () => { |
| 1274 | + peak = Math.max(peak, ++resolving) |
| 1275 | + await delayed(50, undefined) |
| 1276 | + resolving-- |
| 1277 | + return { account, accessToken: 'secret' } |
| 1278 | + }) |
| 1279 | + mocks.search.mockImplementation(() => delayed(3000, { documents: [document] })) |
| 1280 | + let result: Awaited<ReturnType<typeof searchLiveKnowledge.execute>> | undefined |
| 1281 | + const pending = searchLiveKnowledge.execute({ principal, input }).then((value) => { |
| 1282 | + result = value |
| 1283 | + }) |
| 1284 | + await vi.advanceTimersByTimeAsync(3500) |
| 1285 | + expect(result?.live?.accounts.map(({ status }) => status)).toEqual( |
| 1286 | + Array.from({ length: 20 }, () => 'ok') |
| 1287 | + ) |
| 1288 | + expect(peak).toBeLessThanOrEqual(6) |
| 1289 | + await pending |
| 1290 | + }) |
| 1291 | + |
| 1292 | + it('times an account out at its deadline when its session never opens', async () => { |
| 1293 | + mocks.service.mockReturnValue(new Promise(() => {})) |
| 1294 | + let result: Awaited<ReturnType<typeof searchLiveKnowledge.execute>> | undefined |
| 1295 | + void searchLiveKnowledge.execute({ principal, input }).then((value) => { |
| 1296 | + result = value |
| 1297 | + }) |
| 1298 | + await vi.advanceTimersByTimeAsync(12_000) |
| 1299 | + expect(result?.live?.accounts[0]?.status).toBe('timeout') |
| 1300 | + expect(mocks.search).not.toHaveBeenCalled() |
| 1301 | + }) |
| 1302 | + |
| 1303 | + it('gives each account its own deadline from its slot, and skips one the search left no time', async () => { |
| 1304 | + mocks.accounts.mockResolvedValue(accounts(13)) |
| 1305 | + mocks.resolveAccount.mockImplementation(() => |
| 1306 | + delayed(9500, { account, accessToken: 'secret' }) |
| 1307 | + ) |
| 1308 | + mocks.search.mockImplementation(async (_provider, _client, search) => { |
| 1309 | + search.signal?.throwIfAborted() |
| 1310 | + return { documents: [document] } |
| 1311 | + }) |
| 1312 | + const pending = searchLiveKnowledge.execute({ principal, input }) |
| 1313 | + await vi.advanceTimersByTimeAsync(20_000) |
| 1314 | + const statuses = (await pending).live?.accounts ?? [] |
| 1315 | + /** The second wave waits 9.5s for a slot, then still has its full account deadline. */ |
| 1316 | + expect(statuses.slice(0, 12).map(({ status }) => status)).toEqual( |
| 1317 | + Array.from({ length: 12 }, () => 'ok') |
| 1318 | + ) |
| 1319 | + expect(statuses[12]).toMatchObject({ |
| 1320 | + accountId: 'account-12', |
| 1321 | + status: 'timeout', |
| 1322 | + message: expect.stringContaining('Not searched before the search deadline'), |
| 1323 | + }) |
| 1324 | + expect(mocks.resolveAccount).toHaveBeenCalledTimes(12) |
| 1325 | + }) |
| 1326 | + }) |
1156 | 1327 | }) |
0 commit comments