@@ -364,7 +364,7 @@ function createAppServerHarness(
364364 } = { } ,
365365) {
366366 const requests : Array < { method : string ; params : unknown } > = [ ] ;
367- let notify : ( notification : CodexServerNotification ) => Promise < void > = async ( ) => undefined ;
367+ let notifyHandler : ( ( notification : CodexServerNotification ) => Promise < void > ) | undefined ;
368368 let handleServerRequest : AppServerRequestHandler | undefined ;
369369 const closeHandlers = new Set < ( ) => void > ( ) ;
370370 const request = vi . fn ( async ( method : string , params ?: unknown , requestOptions ?: unknown ) => {
@@ -377,9 +377,15 @@ function createAppServerHarness(
377377 return {
378378 getServerVersion : ( ) => "0.132.0" ,
379379 request,
380- addNotificationHandler : ( handler : typeof notify ) => {
381- notify = handler ;
382- return ( ) => undefined ;
380+ addNotificationHandler : (
381+ handler : ( notification : CodexServerNotification ) => Promise < void > ,
382+ ) => {
383+ notifyHandler = handler ;
384+ return ( ) => {
385+ if ( notifyHandler === handler ) {
386+ notifyHandler = undefined ;
387+ }
388+ } ;
383389 } ,
384390 addRequestHandler : ( handler : AppServerRequestHandler ) => {
385391 handleServerRequest = handler ;
@@ -400,6 +406,18 @@ function createAppServerHarness(
400406 return handleServerRequest ! ;
401407 } ;
402408
409+ const waitForNotificationHandler = async ( ) => {
410+ await vi . waitFor ( ( ) => expect ( notifyHandler ) . toBeTypeOf ( "function" ) , {
411+ interval : 1 ,
412+ timeout : appServerHarnessWait . timeout ,
413+ } ) ;
414+ return notifyHandler ! ;
415+ } ;
416+ const sendNotification = async ( notification : CodexServerNotification ) => {
417+ const handler = notifyHandler ?? ( await waitForNotificationHandler ( ) ) ;
418+ await handler ( notification ) ;
419+ } ;
420+
403421 return {
404422 request,
405423 requests,
@@ -419,15 +437,15 @@ function createAppServerHarness(
419437 ) ;
420438 } ,
421439 async notify ( notification : CodexServerNotification ) {
422- await notify ( notification ) ;
440+ await sendNotification ( notification ) ;
423441 } ,
424442 waitForServerRequestHandler,
425443 async handleServerRequest ( request : Parameters < AppServerRequestHandler > [ 0 ] ) {
426444 const handler = await waitForServerRequestHandler ( ) ;
427445 return handler ( request ) ;
428446 } ,
429447 async completeTurn ( params : { threadId : string ; turnId : string } ) {
430- await notify ( {
448+ await sendNotification ( {
431449 method : "turn/completed" ,
432450 params : {
433451 threadId : params . threadId ,
0 commit comments