Disconnect logic and UI fix

handler.py

@@ -96,25 +96,26 @@ class AudioLoop:
                 self.ws = ws
                 await self.startup(tools)
 
-                async with asyncio.TaskGroup() as tg:
-                    send_task = tg.create_task(self.send_realtime())
-                    receive_task = tg.create_task(self.receive_audio())
-                    await asyncio.Future()  # Keep running until canceled
+                try:
+                    async with asyncio.TaskGroup() as tg:
+                        send_task = tg.create_task(self.send_realtime())
+                        receive_task = tg.create_task(self.receive_audio())
+                        await asyncio.Future()  # Keep running until canceled
+                finally:
+                    # Clean up tasks and connection
+                    if 'send_task' in locals():
+                        send_task.cancel()
+                    if 'receive_task' in locals():
+                        receive_task.cancel()                        
+                    try:
+                        await self.ws.close()
+                        print("[AudioLoop] Closed WebSocket connection")
+                    except Exception as e:
+                        print(f"[AudioLoop] Error closing Gemini connection: {e}")
+                    print("[AudioLoop] Cleanup complete")
 
         except asyncio.CancelledError:
-            print("[AudioLoop] Received cancel signal, cleaning up...")
-            # Cancel our tasks if they're still running
-            if 'send_task' in locals():
-                send_task.cancel()
-            if 'receive_task' in locals():
-                receive_task.cancel()
-            # Close Gemini connection if we have one
-            if self.ws and not self.ws.closed:
-                try:
-                    await self.ws.close()
-                except Exception as e:
-                    print(f"[AudioLoop] Error closing Gemini connection: {e}")
-            print("[AudioLoop] Cleanup complete")
+            print("[AudioLoop] Cancelled")
         except Exception as e:
             traceback.print_exc()
             raise

index.html

@@ -117,7 +117,7 @@
       <button id="connectButton" onclick="toggleConnection()">
         <span style="margin-right: 0.3em;">🔌</span> Connect to Server
       </button>
-      <button id="voiceStartButton" class="voice-start" onclick="startCapture()">
+      <button id="voiceStartButton" class="voice-start" onclick="startCapture()" style="display: none;">
         <span style="margin-right: 0.3em;">🎤</span> Start Voice Chat
       </button>
       <button id="voiceStopButton" class="voice-stop" onclick="stopCapture()" style="display: none;">

webapp.py

@@ -102,6 +102,8 @@ async def websocket_endpoint(websocket: WebSocket):
 
         except WebSocketDisconnect:
             print("[from_client_to_gemini] Client disconnected.")
+            #del audio_loop
+            loop_task.cancel()
         except Exception as e:
             print("[from_client_to_gemini] Error:", e)
 
@@ -121,6 +123,7 @@ async def websocket_endpoint(websocket: WebSocket):
 
         except WebSocketDisconnect:
             print("[from_gemini_to_client] Client disconnected.")
+            audio_loop.stop()
         except Exception as e:
             print("[from_gemini_to_client] Error:", e)