first commit

backend/app/routers/collab.py

@@ -0,0 +1,97 @@
+from collections import defaultdict
+
+from fastapi import APIRouter, WebSocket, WebSocketDisconnect
+from sqlalchemy.orm import Session
+
+from app.database import SessionLocal
+from app.mongo import get_mongo_db
+from app.services.document_service import get_document
+from app.services.user_service import get_user_by_email
+from app.utils.security import verify_token
+
+router = APIRouter()
+
+rooms: dict[str, set[WebSocket]] = defaultdict(set)
+presence: dict[str, dict[int, str]] = defaultdict(dict)
+
+
+async def _user_from_token(token: str):
+    subject = verify_token(token, expected_type="access")
+    if not subject:
+        return None
+    db: Session = SessionLocal()
+    try:
+        return get_user_by_email(db, subject)
+    finally:
+        db.close()
+
+
+@router.websocket("/ws/collab/{doc_id}")
+async def collab_socket(websocket: WebSocket, doc_id: str) -> None:
+    token = websocket.query_params.get("token", "")
+    user = await _user_from_token(token)
+    if user is None:
+        await websocket.close(code=4401)
+        return
+
+    try:
+        mongo = get_mongo_db()
+    except RuntimeError:
+        await websocket.close(code=1013)
+        return
+
+    doc = await get_document(mongo, doc_id, user.id)
+    if doc is None:
+        await websocket.close(code=4403)
+        return
+
+    await websocket.accept()
+    rooms[doc_id].add(websocket)
+    presence[doc_id][user.id] = user.full_name or user.username
+
+    async def broadcast(payload: dict, skip: WebSocket | None = None) -> None:
+        dead: list[WebSocket] = []
+        for client in rooms[doc_id]:
+            if client is skip:
+                continue
+            try:
+                await client.send_json(payload)
+            except Exception:
+                dead.append(client)
+        for client in dead:
+            rooms[doc_id].discard(client)
+
+    await broadcast(
+        {
+            "type": "presence",
+            "peers": [{"id": uid, "name": name} for uid, name in presence[doc_id].items()],
+        }
+    )
+
+    try:
+        while True:
+            message = await websocket.receive_json()
+            if message.get("type") == "content":
+                await broadcast(
+                    {
+                        "type": "content",
+                        "content": message.get("content"),
+                        "from": user.id,
+                    },
+                    skip=websocket,
+                )
+    except WebSocketDisconnect:
+        pass
+    finally:
+        rooms[doc_id].discard(websocket)
+        presence[doc_id].pop(user.id, None)
+        if rooms[doc_id]:
+            await broadcast(
+                {
+                    "type": "presence",
+                    "peers": [{"id": uid, "name": name} for uid, name in presence[doc_id].items()],
+                }
+            )
+        else:
+            rooms.pop(doc_id, None)
+            presence.pop(doc_id, None)