@@ -382,12 +382,19 @@ describe("LLM API Tests", () => {
382382 message : "Streaming initiated successfully"
383383 } ) ;
384384
385- // Wait for async streaming operations to complete
386- await new Promise ( resolve => setTimeout ( resolve , 100 ) ) ;
387-
388385 // Import ws service to access mock
389386 const ws = ( await import ( "../../services/ws.js" ) ) . default ;
390387
388+ // Wait for async streaming operations to complete
389+ await vi . waitFor ( ( ) => {
390+ expect ( ws . sendMessageToAllClients ) . toHaveBeenCalledWith ( {
391+ type : 'llm-stream' ,
392+ chatNoteId : testChatId ,
393+ content : ' world!' ,
394+ done : true
395+ } ) ;
396+ } , { timeout : 1000 , interval : 50 } ) ;
397+
391398 // Verify WebSocket messages were sent
392399 expect ( ws . sendMessageToAllClients ) . toHaveBeenCalledWith ( {
393400 type : 'llm-stream' ,
@@ -539,12 +546,19 @@ describe("LLM API Tests", () => {
539546
540547 expect ( response . status ) . toBe ( 200 ) ;
541548
542- // Wait for async streaming operations to complete
543- await new Promise ( resolve => setTimeout ( resolve , 100 ) ) ;
544-
545549 // Import ws service to access mock
546550 const ws = ( await import ( "../../services/ws.js" ) ) . default ;
547551
552+ // Wait for async streaming operations to complete
553+ await vi . waitFor ( ( ) => {
554+ expect ( ws . sendMessageToAllClients ) . toHaveBeenCalledWith ( {
555+ type : 'llm-stream' ,
556+ chatNoteId : testChatId ,
557+ thinking : 'Formulating response...' ,
558+ done : false
559+ } ) ;
560+ } , { timeout : 1000 , interval : 50 } ) ;
561+
548562 // Verify thinking messages
549563 expect ( ws . sendMessageToAllClients ) . toHaveBeenCalledWith ( {
550564 type : 'llm-stream' ,
@@ -589,12 +603,26 @@ describe("LLM API Tests", () => {
589603
590604 expect ( response . status ) . toBe ( 200 ) ;
591605
592- // Wait for async streaming operations to complete
593- await new Promise ( resolve => setTimeout ( resolve , 100 ) ) ;
594-
595606 // Import ws service to access mock
596607 const ws = ( await import ( "../../services/ws.js" ) ) . default ;
597608
609+ // Wait for async streaming operations to complete
610+ await vi . waitFor ( ( ) => {
611+ expect ( ws . sendMessageToAllClients ) . toHaveBeenCalledWith ( {
612+ type : 'llm-stream' ,
613+ chatNoteId : testChatId ,
614+ toolExecution : {
615+ tool : 'calculator' ,
616+ args : { expression : '2 + 2' } ,
617+ result : '4' ,
618+ toolCallId : 'call_123' ,
619+ action : 'execute' ,
620+ error : undefined
621+ } ,
622+ done : false
623+ } ) ;
624+ } , { timeout : 1000 , interval : 50 } ) ;
625+
598626 // Verify tool execution message
599627 expect ( ws . sendMessageToAllClients ) . toHaveBeenCalledWith ( {
600628 type : 'llm-stream' ,
0 commit comments