Spaces:
Configuration error
Configuration error
File size: 7,168 Bytes
78013c4 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 | import asyncio
import json
from typing import Dict, Set
from fastapi import APIRouter, WebSocket, WebSocketDisconnect
from app.api.dependencies import get_current_user_from_token
from app.core.logging_config import get_logger
logger = get_logger("app.websocket")
router = APIRouter(prefix="/ws", tags=["WebSocket"])
_PING_INTERVAL = 25
_PONG_TIMEOUT = 15
class ConnectionManager:
def __init__(self):
self.student_connections: Dict[str, Set[WebSocket]] = {}
self.teacher_connections: Dict[str, Set[WebSocket]] = {}
self._heartbeat_task: asyncio.Task | None = None
async def connect_student(self, websocket: WebSocket, student_id: str):
await websocket.accept()
self.student_connections.setdefault(student_id, set()).add(websocket)
logger.info("WebSocket student connected: %s", student_id)
async def connect_teacher(self, websocket: WebSocket, teacher_id: str):
await websocket.accept()
self.teacher_connections.setdefault(teacher_id, set()).add(websocket)
logger.info("WebSocket teacher connected: %s", teacher_id)
def disconnect(self, websocket: WebSocket, user_type: str, user_id: str):
connections = self.student_connections if user_type == "student" else self.teacher_connections
if user_id in connections:
connections[user_id].discard(websocket)
if not connections[user_id]:
del connections[user_id]
logger.info("WebSocket %s disconnected: %s", user_type, user_id)
async def send_personal_message(self, message: dict, student_id: str):
conns = self.student_connections.get(student_id)
if not conns:
return
disconnected = set()
for connection in conns:
try:
await connection.send_json(message)
except Exception as e:
logger.warning("Failed to send message to %s: %s", student_id, e)
disconnected.add(connection)
for conn in disconnected:
self.student_connections[student_id].discard(conn)
async def broadcast_to_teachers(self, message: dict):
disconnected = set()
for teacher_id, conns in self.teacher_connections.items():
for conn in conns:
try:
await conn.send_json(message)
except Exception as e:
logger.warning("Failed to send to teacher %s: %s", teacher_id, e)
disconnected.add(conn)
for conn in disconnected:
for teacher_id, conns in self.teacher_connections.items():
conns.discard(conn)
if not conns:
del self.teacher_connections[teacher_id]
async def _heartbeat_loop(self):
while True:
await asyncio.sleep(_PING_INTERVAL)
ping = {"type": "ping"}
disconnected = set()
for sid, conns in list(self.student_connections.items()):
for conn in list(conns):
try:
await asyncio.wait_for(
conn.send_json(ping), timeout=_PONG_TIMEOUT
)
except Exception:
disconnected.add((conn, "student", sid))
for tid, conns in list(self.teacher_connections.items()):
for conn in list(conns):
try:
await asyncio.wait_for(
conn.send_json(ping), timeout=_PONG_TIMEOUT
)
except Exception:
disconnected.add((conn, "teacher", tid))
for conn, utype, uid in disconnected:
self.disconnect(conn, utype, uid)
if disconnected:
logger.info(
"Heartbeat cleaned %d stale connections", len(disconnected)
)
def start_heartbeat(self):
if self._heartbeat_task is None:
self._heartbeat_task = asyncio.create_task(self._heartbeat_loop())
logger.info("WebSocket heartbeat started")
def stop_heartbeat(self):
if self._heartbeat_task is not None:
self._heartbeat_task.cancel()
self._heartbeat_task = None
logger.info("WebSocket heartbeat stopped")
@property
def total_connections(self) -> int:
student_count = sum(len(c) for c in self.student_connections.values())
teacher_count = sum(len(c) for c in self.teacher_connections.values())
return student_count + teacher_count
manager = ConnectionManager()
@router.websocket("/connect")
async def websocket_endpoint(websocket: WebSocket):
try:
await websocket.accept()
auth_data = await websocket.receive_text()
auth_json = json.loads(auth_data)
if auth_json.get("type") != "auth" or not auth_json.get("token"):
await websocket.close(code=1008, reason="Authentication required")
return
user = await get_current_user_from_token(auth_json["token"])
if not user:
await websocket.close(code=1008, reason="Unauthorized")
return
if user.role == "STUDENT" and user.student:
student_id = user.student.id
await manager.connect_student(websocket, student_id)
await websocket.send_json({"type": "connected", "message": "WebSocket connection established", "user_id": student_id, "role": "student"})
try:
while True:
data = await asyncio.wait_for(websocket.receive_text(), timeout=_PING_INTERVAL)
if data == "ping":
await websocket.send_json({"type": "pong"})
except asyncio.TimeoutError:
logger.info("WebSocket ping timeout for student %s", student_id)
except WebSocketDisconnect:
manager.disconnect(websocket, "student", student_id)
elif user.role == "TEACHER" and user.teacher:
teacher_id = user.teacher.id
await manager.connect_teacher(websocket, teacher_id)
await websocket.send_json({"type": "connected", "message": "WebSocket connection established", "user_id": teacher_id, "role": "teacher"})
try:
while True:
data = await asyncio.wait_for(websocket.receive_text(), timeout=_PING_INTERVAL)
if data == "ping":
await websocket.send_json({"type": "pong"})
except asyncio.TimeoutError:
logger.info("WebSocket ping timeout for teacher %s", teacher_id)
except WebSocketDisconnect:
manager.disconnect(websocket, "teacher", teacher_id)
else:
await websocket.close(code=1008, reason="Unauthorized: Student or Teacher profile required")
except Exception as e:
logger.error("WebSocket error: %s", e, exc_info=True)
try:
await websocket.close(code=1011, reason="Internal server error")
except Exception:
pass
|