Upd ws foward_progress

app/routers/ws_progress.py

@@ -26,28 +26,41 @@ async def _delete_everything(doc_id: str):
 
 async def forward_progress(websocket: WebSocket, document_id: str):
     """Handle state change allowing frontend to connect and update seamlessly."""
-    db = get_db()
+    logger.info(f"📡 WebSocket accepted for doc {document_id}")
     try:
+        db = get_db()
+        doc = await db.documents.find_one({"_id": document_id})
+        if not doc:
+            await websocket.send_json({"status": "NOT_FOUND"})
+            return
+        # Valid doc
         while True:
             doc = await db.documents.find_one({"_id": document_id})
             if not doc:
                 await websocket.send_json({"status": "NOT_FOUND"})
-            else:
-                status = doc.get("status", "UNKNOWN")
-                await websocket.send_json({"status": status})
-                if status in {"READY", "FAILED"}:
-                    break
-            await asyncio.sleep(2)
-
-    except WebSocketDisconnect:
-        logger.info(f"💨 client closed WS for {document_id} – deleting artefacts")
-        await _delete_everything(document_id)
-
+                return
+            # Get status real-time
+            status = doc.get("status")
+            if status == "READY":
+                await websocket.send_json({
+                    "status": "READY",
+                    "id": doc["_id"],
+                    "title": doc.get("title"),
+                    "source": doc.get("metadata", {}).get("source", "unknown"),
+                    "documentId": doc["_id"],
+                    "uri": f"/import/textbook/{doc['_id']}",
+                })
+                break
+            elif status == "FAILED":
+                await websocket.send_json({"status": "FAILED"})
+                break
+            await asyncio.sleep(1.5)
     except Exception as e:
-        logger.exception(f"📡 WS loop crashed for {document_id}: {e}")
-        with contextlib.suppress(Exception):
+        logger.exception(f"📡 WebSocket failed for doc {document_id}: {e}")
+        try:
             await websocket.send_json({"status": "ERROR"})
-
-    finally:
-        with contextlib.suppress(Exception):
             await websocket.close()
+        except Exception:
+            pass
+    finally:
+        logger.info(f"📡 WebSocket closed for doc {document_id}")