Bugs are fixed in the backend, it works now

2026-08-01 08:40:53 +00:00
parent cd56637d2a
commit b6ace3f68b
8 changed files with 180 additions and 10 deletions
@@ -0,0 +1,46 @@
+"""WebSocket endpoint for real‑time progress updates."""
+
+import json
+
+from fastapi import APIRouter, WebSocket, WebSocketDisconnect
+
+from backend.services.task_manager import task_manager
+from backend.services.progress_publisher import register_connection, unregister_connection
+
+router = APIRouter(tags=["websocket"])
+
+
+@router.websocket("/ws/{task_id}")
+async def websocket_endpoint(websocket: WebSocket, task_id: str):
+    await websocket.accept()
+
+    register_connection(task_id, websocket)
+
+    try:
+        # Send initial status
+        status = task_manager.get_status(task_id)
+        await websocket.send_json({
+            "type": "status",
+            "data": status
+        })
+
+        while True:
+            data = await websocket.receive_text()
+            try:
+                message = json.loads(data)
+                if message.get("type") == "ping":
+                    await websocket.send_json({"type": "pong"})
+                elif message.get("type") == "get_status":
+                    status = task_manager.get_status(task_id)
+                    await websocket.send_json({
+                        "type": "status",
+                        "data": status
+                    })
+            except json.JSONDecodeError:
+                pass
+
+    except WebSocketDisconnect:
+        unregister_connection(task_id, websocket)
+    except Exception as e:
+        unregister_connection(task_id, websocket)
+        print(f"WebSocket error: {e}")