11import { ErrorCode , McpError } from '@modelcontextprotocol/sdk/types.js'
2- import { dbChainMockFns , resetDbChainMock } from '@sim/testing'
2+ import { createDeferred , dbChainMockFns , resetDbChainMock } from '@sim/testing'
33import { createSessionPrincipal } from '@sim/testing/factories/principal.factory'
44import { setEnv } from '@sim/testing/mocks/env.mock'
55import {
@@ -11,7 +11,7 @@ import {
1111 knowledgeContextsMockFns ,
1212} from '@sim/testing/mocks/knowledge-contexts.mock'
1313import { workspaceAuthzMock , workspaceAuthzMockFns } from '@sim/testing/mocks/workspace-authz.mock'
14- import { beforeEach , describe , expect , it , vi } from 'vitest'
14+ import { afterEach , beforeEach , describe , expect , it , vi } from 'vitest'
1515
1616const mocks = vi . hoisted ( ( ) => ( {
1717 accounts : vi . fn ( ) ,
@@ -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 ( ) => {
@@ -1105,4 +1148,94 @@ describe('authorized live retrieval', () => {
11051148 } )
11061149 ) . toBe ( false )
11071150 } )
1151+
1152+ const accounts = ( count : number ) =>
1153+ Array . from ( { length : count } , ( _ , index ) => ( {
1154+ ...account ,
1155+ id : `account-${ String ( index ) . padStart ( 2 , '0' ) } ` ,
1156+ displayName : `Drive ${ String ( index ) . padStart ( 2 , '0' ) } ` ,
1157+ } ) )
1158+
1159+ it ( 'reports each account beyond the search limit instead of omitting it' , async ( ) => {
1160+ mocks . accounts . mockResolvedValue ( accounts ( 21 ) )
1161+ const result = await searchLiveKnowledge . execute ( { principal, input } )
1162+ const statuses = result . live ?. accounts ?? [ ]
1163+ expect ( statuses ) . toHaveLength ( 21 )
1164+ expect ( new Set ( statuses . map ( ( { accountId } ) => accountId ) ) . size ) . toBe ( 21 )
1165+ expect ( statuses . at ( - 1 ) ) . toMatchObject ( {
1166+ accountId : 'account-20' ,
1167+ status : 'unavailable' ,
1168+ message : expect . stringContaining ( 'Not searched' ) ,
1169+ } )
1170+ expect ( result . retrieval . status ) . toBe ( 'partial' )
1171+ } )
1172+
1173+ describe ( 'latency' , ( ) => {
1174+ beforeEach ( ( ) => {
1175+ vi . useFakeTimers ( )
1176+ vi . spyOn ( AbortSignal , 'timeout' ) . mockImplementation ( ( milliseconds ) => {
1177+ const controller = new AbortController ( )
1178+ setTimeout (
1179+ ( ) => controller . abort ( new DOMException ( 'deadline' , 'TimeoutError' ) ) ,
1180+ milliseconds
1181+ )
1182+ return controller . signal
1183+ } )
1184+ } )
1185+ afterEach ( ( ) => {
1186+ vi . mocked ( AbortSignal . timeout ) . mockRestore ( )
1187+ vi . useRealTimers ( )
1188+ } )
1189+
1190+ /** Settles after `ms` of fake time, like a provider round trip. */
1191+ const delayed = < T > ( ms : number , value : T ) =>
1192+ new Promise < T > ( ( resolve ) => setTimeout ( ( ) => resolve ( value ) , ms ) )
1193+
1194+ it ( 'searches every account at once while bounding concurrent account resolution' , async ( ) => {
1195+ mocks . accounts . mockResolvedValue ( accounts ( 20 ) )
1196+ let resolving = 0
1197+ let peak = 0
1198+ mocks . resolveAccount . mockImplementation ( async ( ) => {
1199+ peak = Math . max ( peak , ++ resolving )
1200+ await delayed ( 50 , undefined )
1201+ resolving --
1202+ return { account, accessToken : 'secret' }
1203+ } )
1204+ mocks . search . mockImplementation ( ( ) => delayed ( 3000 , { documents : [ document ] } ) )
1205+ let result : Awaited < ReturnType < typeof searchLiveKnowledge . execute > > | undefined
1206+ const pending = searchLiveKnowledge . execute ( { principal, input } ) . then ( ( value ) => {
1207+ result = value
1208+ } )
1209+ await vi . advanceTimersByTimeAsync ( 3500 )
1210+ expect ( result ?. live ?. accounts . map ( ( { status } ) => status ) ) . toEqual (
1211+ Array . from ( { length : 20 } , ( ) => 'ok' )
1212+ )
1213+ expect ( peak ) . toBeLessThanOrEqual ( 6 )
1214+ await pending
1215+ } )
1216+
1217+ it ( 'gives each account its own deadline from its slot, and skips one the search left no time' , async ( ) => {
1218+ mocks . accounts . mockResolvedValue ( accounts ( 13 ) )
1219+ mocks . resolveAccount . mockImplementation ( ( ) =>
1220+ delayed ( 9500 , { account, accessToken : 'secret' } )
1221+ )
1222+ mocks . search . mockImplementation ( async ( _provider , _client , search ) => {
1223+ search . signal ?. throwIfAborted ( )
1224+ return { documents : [ document ] }
1225+ } )
1226+ const pending = searchLiveKnowledge . execute ( { principal, input } )
1227+ await vi . advanceTimersByTimeAsync ( 20_000 )
1228+ const statuses = ( await pending ) . live ?. accounts ?? [ ]
1229+ /** The second wave waits 9.5s for a slot, then still has its full account deadline. */
1230+ expect ( statuses . slice ( 0 , 12 ) . map ( ( { status } ) => status ) ) . toEqual (
1231+ Array . from ( { length : 12 } , ( ) => 'ok' )
1232+ )
1233+ expect ( statuses [ 12 ] ) . toMatchObject ( {
1234+ accountId : 'account-12' ,
1235+ status : 'timeout' ,
1236+ message : expect . stringContaining ( 'Not searched before the search deadline' ) ,
1237+ } )
1238+ expect ( mocks . resolveAccount ) . toHaveBeenCalledTimes ( 12 )
1239+ } )
1240+ } )
11081241} )
0 commit comments