worked backup

backend/src/app.py

@@ -50,25 +50,41 @@ app.include_router(settings.router, prefix="/api/settings", tags=["Settings"])
 # @SEMANTICS: websocket, logs, streaming, real-time
 # @PURPOSE: Provides a WebSocket endpoint for clients to connect to and receive real-time log entries for a specific task.
 @app.websocket("/ws/logs/{task_id}")
-async def websocket_endpoint(websocket: WebSocket, task_id: str, task_manager=Depends(get_task_manager)):
+async def websocket_endpoint(websocket: WebSocket, task_id: str):
     await websocket.accept()
-    logger.info(f"WebSocket connection established for task {task_id}")
+    logger.info(f"WebSocket connection accepted for task {task_id}")
+    task_manager = get_task_manager()
+    queue = await task_manager.subscribe_logs(task_id)
     try:
         # Send initial logs if any
         initial_logs = task_manager.get_task_logs(task_id)
         for log_entry in initial_logs:
-            await websocket.send_json(log_entry.dict())
+            # Convert datetime to string for JSON serialization
+            log_dict = log_entry.dict()
+            log_dict['timestamp'] = log_dict['timestamp'].isoformat()
+            await websocket.send_json(log_dict)
 
-        # Keep connection alive, ideally stream new logs as they come
-        # This part requires a more sophisticated log streaming mechanism (e.g., queues, pub/sub)
-        # For now, it will just keep the connection open and send initial logs.
+        # Stream new logs
+        logger.info(f"Starting log stream for task {task_id}")
         while True:
-            await asyncio.sleep(1) # Keep connection alive, send heartbeat or check for new logs
-            # In a real system, new logs would be pushed here
+            log_entry = await queue.get()
+            log_dict = log_entry.dict()
+            log_dict['timestamp'] = log_dict['timestamp'].isoformat()
+            await websocket.send_json(log_dict)
+            
+            # If task is finished, we could potentially close the connection
+            # but let's keep it open for a bit or until the client disconnects
+            if "Task completed successfully" in log_entry.message or "Task failed" in log_entry.message:
+                # Wait a bit to ensure client receives the last message
+                await asyncio.sleep(2)
+                break
+
     except WebSocketDisconnect:
         logger.info(f"WebSocket connection disconnected for task {task_id}")
     except Exception as e:
         logger.error(f"WebSocket error for task {task_id}: {e}")
+    finally:
+        task_manager.unsubscribe_logs(task_id, queue)
 
 # [/DEF]