Hoàn thành chức năng collaborative

backend/app/routers/ws.py

@@ -0,0 +1,172 @@
+"""
+WebSocket collaboration endpoint.
+
+URL: ws://.../ws/pdf/{pdf_id}?token=<JWT>
+
+Each PDF has its own "room". When a client connects it broadcasts a `presence`
+event to the room. All annotation mutations are forwarded verbatim to every
+other client in the room.
+
+Message schema (JSON):
+  → client sends:
+      { "type": "object_add",    "payload": <fabric object JSON> }
+      { "type": "object_remove", "payload": { "obj_id": "<id>" } }
+      { "type": "clear" }
+      { "type": "cursor",        "payload": { "x": 0, "y": 0, "page": 1 } }
+      { "type": "ping" }
+
+  ← server sends to ALL others in room:
+      same messages, with "user_id" / "username" / "color" injected
+
+  ← server sends to ALL (including sender) on join/leave:
+      { "type": "presence", "users": [ { "user_id", "username", "color" } ] }
+"""
+
+import asyncio
+import json
+import logging
+from typing import Any
+
+from fastapi import APIRouter, WebSocket, WebSocketDisconnect, status
+from jose import JWTError
+
+from ..security import decode_access_token
+from ..database import get_db
+from ..models import User
+
+logger = logging.getLogger(__name__)
+router = APIRouter()
+
+# ── Deterministic colour per user (hue based on user_id) ─────────────────────
+
+def _user_color(user_id: int) -> str:
+    hue = (user_id * 61) % 360          # spread nicely around the wheel
+    return f"hsl({hue},70%,50%)"
+
+
+# ── Room manager ──────────────────────────────────────────────────────────────
+
+class _Room:
+    def __init__(self) -> None:
+        # websocket → user info dict
+        self._clients: dict[WebSocket, dict[str, Any]] = {}
+        self._lock = asyncio.Lock()
+
+    async def join(self, ws: WebSocket, user_info: dict[str, Any]) -> None:
+        async with self._lock:
+            self._clients[ws] = user_info
+        await self._broadcast_presence()
+
+    async def leave(self, ws: WebSocket) -> None:
+        async with self._lock:
+            self._clients.pop(ws, None)
+        await self._broadcast_presence()
+
+    @property
+    def user_list(self) -> list[dict[str, Any]]:
+        """Deduplicated by user_id — same user may have multiple WS connections."""
+        seen: set[int] = set()
+        result: list[dict[str, Any]] = []
+        for info in self._clients.values():
+            uid = info["user_id"]
+            if uid not in seen:
+                seen.add(uid)
+                result.append(info)
+        return result
+
+    async def broadcast(self, message: dict[str, Any], exclude: WebSocket | None = None) -> None:
+        """Send message to every client except `exclude`."""
+        async with self._lock:
+            targets = [ws for ws in self._clients if ws is not exclude]
+        for ws in targets:
+            try:
+                await ws.send_json(message)
+            except Exception:
+                pass
+
+    async def broadcast_all(self, message: dict[str, Any]) -> None:
+        """Send message to every client including sender."""
+        await self.broadcast(message, exclude=None)
+
+    async def _broadcast_presence(self) -> None:
+        await self.broadcast_all({"type": "presence", "users": self.user_list})
+
+
+class _RoomManager:
+    def __init__(self) -> None:
+        self._rooms: dict[int, _Room] = {}
+        self._lock = asyncio.Lock()
+
+    async def get_or_create(self, pdf_id: int) -> _Room:
+        async with self._lock:
+            if pdf_id not in self._rooms:
+                self._rooms[pdf_id] = _Room()
+            return self._rooms[pdf_id]
+
+    async def cleanup(self, pdf_id: int) -> None:
+        async with self._lock:
+            room = self._rooms.get(pdf_id)
+            if room and not room._clients:
+                del self._rooms[pdf_id]
+
+
+manager = _RoomManager()
+
+# ── WebSocket endpoint ────────────────────────────────────────────────────────
+
+@router.websocket("/ws/pdf/{pdf_id}")
+async def pdf_collaboration(websocket: WebSocket, pdf_id: int):
+    # ── Auth: JWT passed as query param (cookies aren't reliably sent on WS) ──
+    token = websocket.query_params.get("token")
+    if not token:
+        await websocket.close(code=status.WS_1008_POLICY_VIOLATION)
+        return
+
+    try:
+        user_id = decode_access_token(token)
+    except JWTError:
+        await websocket.close(code=status.WS_1008_POLICY_VIOLATION)
+        return
+
+    # ── Fetch username from DB ────────────────────────────────────────────────
+    db = next(get_db())
+    try:
+        user: User | None = db.get(User, user_id)
+    finally:
+        db.close()
+
+    if user is None:
+        await websocket.close(code=status.WS_1008_POLICY_VIOLATION)
+        return
+
+    color = _user_color(user_id)
+    user_info = {"user_id": user_id, "username": user.username, "color": color}
+
+    await websocket.accept()
+    room = await manager.get_or_create(pdf_id)
+    await room.join(websocket, user_info)
+
+    try:
+        while True:
+            raw = await websocket.receive_text()
+            try:
+                msg = json.loads(raw)
+            except json.JSONDecodeError:
+                continue
+
+            msg_type = msg.get("type")
+
+            if msg_type == "ping":
+                await websocket.send_json({"type": "pong"})
+                continue
+
+            # Inject sender identity and forward to all other clients
+            if msg_type in {"object_add", "object_remove", "clear", "cursor"}:
+                msg.update(user_info)
+                await room.broadcast(msg, exclude=websocket)
+
+    except WebSocketDisconnect:
+        pass
+    finally:
+        await room.leave(websocket)
+        await manager.cleanup(pdf_id)