Fix test client to wait for TURN_COMPLETE instead of fixed timeout

Author: oadultradeepfield

scripts/test_client.py

@@ -126,6 +126,7 @@ async def handle_server_message(
     output_dir: Path | None,
     verbose: bool,
     audio_buffer: list[bytes],
+    turn_complete_event: asyncio.Event | None = None,
 ) -> None:
     """Process a message from the server."""
     if isinstance(message, bytes):
@@ -160,6 +161,8 @@ async def handle_server_message(
                 path = save_audio(combined, output_dir)
                 logger.info("Saved audio: %s (%d bytes)", path, len(combined))
                 audio_buffer.clear()
+            if turn_complete_event:
+                turn_complete_event.set()
         case "INTERRUPTED":
             logger.info("Agent interrupted")
             audio_buffer.clear()
@@ -212,12 +215,13 @@ async def receive_messages(
     ws: ClientConnection,
     output_dir: Path | None,
     verbose: bool,
+    turn_complete_event: asyncio.Event,
 ) -> None:
     """Receive and handle all messages from the server."""
     audio_buffer: list[bytes] = []
 
     async for message in ws:
-        await handle_server_message(message, output_dir, verbose, audio_buffer)
+        await handle_server_message(message, output_dir, verbose, audio_buffer, turn_complete_event)
 
 
 async def run_client(
@@ -262,12 +266,18 @@ async def run_client(
             logger.info("Sent IMAGE frame")
 
         ping_task = asyncio.create_task(send_ping_periodically(ws))
-        receive_task = asyncio.create_task(receive_messages(ws, output_dir, verbose))
+        turn_complete_event = asyncio.Event()
+        receive_task = asyncio.create_task(
+            receive_messages(ws, output_dir, verbose, turn_complete_event)
+        )
 
         await stream_audio(ws, audio_data, verbose)
 
         logger.info("Waiting for agent response...")
-        await asyncio.sleep(2)
+        try:
+            await asyncio.wait_for(turn_complete_event.wait(), timeout=60.0)
+        except TimeoutError:
+            logger.warning("Timeout waiting for TURN_COMPLETE")
 
         ping_task.cancel()
         with contextlib.suppress(asyncio.CancelledError):