Update app.py

app.py

@@ -198,25 +198,23 @@ threading.Thread(target=ai_manager.modify_research_tasks, daemon=True).start()
 @app.websocket("/ws")
 async def websocket_endpoint(websocket: WebSocket):
     await websocket.accept()
-    print("WebSocket connection accepted")  # Add explicit logging
-    
+    print("WebSocket connection accepted")
     ai_manager.clients.add(websocket)
+    
     try:
-        # Send initial welcome message
+        # 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()
-            print(f"Received WebSocket message: {data}")  # Explicit receive logging
+            print(f"Received message: {data}")
             
-            # Process message through AI manager
-            ai_manager.receive_message("human", data)
+            # Directly broadcast the message
+            for client in ai_manager.clients:
+                await client.send_text(f"Received: {data}")
             
-            # Ensure broadcast is working
-            await ai_manager.broadcast(f"Received: {data}")
     except Exception as e:
-        print(f"WebSocket connection error: {e}")
+        print(f"WebSocket error: {e}")
         ai_manager.clients.remove(websocket)
 
 @app.get("/")