|
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,49 @@ 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('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 | + }) |
743 | 786 | }) |
744 | 787 |
|
745 | 788 | it('gives provider searches the account session cancellation signal', async () => { |
@@ -1153,4 +1196,94 @@ describe('authorized live retrieval', () => { |
1153 | 1196 | }) |
1154 | 1197 | ).toBe(false) |
1155 | 1198 | }) |
| 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 | + }) |
1156 | 1289 | }) |
0 commit comments