Update app.py

app.py

@@ -203,18 +203,26 @@ threading.Thread(target=ai_manager.modify_research_tasks, daemon=True).start()
 
 @app.websocket("/ws")
 async def websocket_endpoint(websocket: WebSocket):
-    """WebSocket connection handler."""
     await websocket.accept()
-    await websocket.send_json({"msg": "Hello WebSocket"})
-
+    print("WebSocket connection accepted")  # Add explicit logging
+    
     ai_manager.clients.add(websocket)
     try:
+        # Send initial welcome message
+        await websocket.send_text("Connected to AI State Machine")
+        print("Sent initial welcome message")  # Debug logging
+        
         while True:
-            data = await websocket.receive_text()  # Receive messages from the user.
-            print(f"Received: {data}")  # Debug: Print the received data.
-            ai_manager.receive_message("human", data)  # Send it to the AI manager.
+            data = await websocket.receive_text()
+            print(f"Received WebSocket message: {data}")  # Explicit receive logging
+            
+            # Process message through AI manager
+            ai_manager.receive_message("human", data)
+            
+            # Ensure broadcast is working
+            await ai_manager.broadcast(f"Received: {data}")
     except Exception as e:
-        print(f"Connection closed with error: {e}")
+        print(f"WebSocket connection error: {e}")
         ai_manager.clients.remove(websocket)
 
 @app.get("/")