app.py improved run performance still some bugs

app.py

@@ -31,7 +31,6 @@ gemini_connected = False
 
 # --- Helper Functions ---
 def encode_audio(data: np.ndarray) -> dict:
-    """Encode Audio data."""
     if not isinstance(data, np.ndarray):
         raise TypeError("encode_audio expected a numpy.ndarray")
     try:
@@ -41,7 +40,6 @@ def encode_audio(data: np.ndarray) -> dict:
         raise
 
 def encode_image(data: np.ndarray) -> dict:
-    """Encode Image data."""
     if not isinstance(data, np.ndarray):
         raise TypeError("encode_image expected a numpy.ndarray")
     try:
@@ -140,9 +138,9 @@ class GeminiHandler(AsyncAudioVideoStreamHandler):
             try:
                 client = genai.Client(api_key=api_key, http_options={"api_version": "v1alpha"})
                 config = {"response_modalities": ["AUDIO"]}
-                async with client.aio.live.connect(  # Use async with, like the original
+                async with client.aio.live.connect(
                     model="gemini-2.0-flash-exp", config=config
-                ) as session:  # <--- Get session from context manager
+                ) as session:
                     self.session = session
                     gemini_connected = True
                     asyncio.create_task(self.receive_audio())
@@ -163,19 +161,22 @@ class GeminiHandler(AsyncAudioVideoStreamHandler):
 
         while not self.quit.is_set():
             try:
-                turn = self.session.receive()  # NO await here, like the original
+                await asyncio.sleep(0)  # Yield to the event loop
+                if self.quit.is_set():
+                    break
+                turn = self.session.receive()
                 async for response in turn:
+                    if self.quit.is_set():
+                        break # Exit inner loop if quit is set.
                     if data := response.data:
                         yield data
             except Exception as e:
                 logger.error(f"Error receiving from Gemini: {e}")
-                # NOTE: Do not exit loop.
-                # The user may need to say something else
-                # if there is a problem.
-                # break
+                self.quit.set() # set quit if we error.
+                break
+
     async def receive_audio(self):
         try:
-            # Correctly use the async generator
             async for audio_response in async_aggregate_bytes_to_16bit(self.generator()):
                 self.audio_queue.put_nowait(audio_response)
         except Exception as e:
@@ -212,12 +213,16 @@ class GeminiHandler(AsyncAudioVideoStreamHandler):
         global gemini_connected
         gemini_connected = False
         logger.info("Shutting down GeminiHandler.")
-        self.quit.set()
+        if self.session:
+            try:
+                #  await self.session.close()  # There is no async close
+                pass
+            except Exception:
+                pass
+        self.quit.set()  # Set quit *after* attempting to close the session
         self.connection = None
         self.args_set.clear()
-        if self.session:
-            # No good async close.
-            pass
+
         self.quit.clear()
         update_gemini_status_sync()