add websockets suppor98

backend/app/main.py

@@ -37,8 +37,13 @@ app.include_router(admin.router, prefix="/api/admin", tags=["admin"])
 app.include_router(statistics.router, prefix="/api/statistics", tags=["statistics"])
 
 # Добавляем WebSocket маршруты
-app.websocket("/api/ws/admin")(notification_manager.admin_endpoint)
-app.websocket("/api/ws/employee/{employee_id}")(notification_manager.employee_endpoint)
+@app.websocket("/api/ws/admin")
+async def admin_websocket(websocket):
+    await notification_manager.admin_endpoint(websocket)
+
+@app.websocket("/api/ws/employee/{employee_id}")
+async def employee_websocket(websocket, employee_id: int):
+    await notification_manager.employee_endpoint(websocket, employee_id)
 
 @app.on_event("startup")
 async def startup_event():

backend/app/websockets/notifications.py

@@ -1,81 +1,93 @@
+"""WebSocket notifications manager"""
 from fastapi import WebSocket
-from typing import Dict, List
-import json
+from typing import Dict, List, Set
 import logging
+import json
 
 logger = logging.getLogger(__name__)
 
 class NotificationManager:
+    """Менеджер WebSocket подключений и уведомлений"""
     def __init__(self):
-        self.active_connections: Dict[str, List[WebSocket]] = {
-            "admin": [],  # Подключения админов
-            "employee": []  # Подключения сотрудников
-        }
-
-    async def connect(self, websocket: WebSocket, client_type: str):
+        self.admin_connections: Set[WebSocket] = set()
+        self.employee_connections: Dict[int, WebSocket] = {}
+        
+    async def connect(self, websocket: WebSocket, client_type: str, employee_id: int = None):
+        """Установка WebSocket соединения"""
         await websocket.accept()
-        self.active_connections[client_type].append(websocket)
-        logger.info(f"New {client_type} connection. Total connections: {len(self.active_connections[client_type])}")
+        if client_type == "admin":
+            self.admin_connections.add(websocket)
+            logger.info("Admin connected to WebSocket")
+        elif client_type == "employee" and employee_id:
+            self.employee_connections[employee_id] = websocket
+            logger.info(f"Employee {employee_id} connected to WebSocket")
 
-    def disconnect(self, websocket: WebSocket, client_type: str):
-        if websocket in self.active_connections[client_type]:
-            self.active_connections[client_type].remove(websocket)
-            logger.info(f"{client_type} disconnected. Remaining connections: {len(self.active_connections[client_type])}")
+    def disconnect(self, websocket: WebSocket, client_type: str, employee_id: int = None):
+        """Закрытие WebSocket соединения"""
+        if client_type == "admin":
+            self.admin_connections.discard(websocket)
+            logger.info("Admin disconnected from WebSocket")
+        elif client_type == "employee" and employee_id:
+            self.employee_connections.pop(employee_id, None)
+            logger.info(f"Employee {employee_id} disconnected from WebSocket")
 
     async def broadcast_to_admins(self, message: dict):
-        """Отправка сообщения всем подключенным админам"""
-        logger.info(f"Broadcasting to admins: {message}")
-        logger.info(f"Number of admin connections: {len(self.active_connections['admin'])}")
-        
-        if not self.active_connections["admin"]:
-            logger.warning("No admin connections available")
-            return
-            
-        disconnected = []
-        
-        for connection in self.active_connections["admin"]:
+        """Отправка сообщения всем админам"""
+        disconnected = set()
+        for websocket in self.admin_connections:
             try:
-                await connection.send_json(message)
-                logger.info("Message sent successfully to admin")
+                await websocket.send_json(message)
             except Exception as e:
                 logger.error(f"Error sending message to admin: {e}")
-                disconnected.append(connection)
-                continue
+                disconnected.add(websocket)
         
         # Удаляем отключенные соединения
-        for connection in disconnected:
-            self.disconnect(connection, "admin")
+        for websocket in disconnected:
+            self.disconnect(websocket, "admin")
 
-    async def broadcast_to_employees(self, employee_id: int, message: dict):
+    async def send_to_employee(self, employee_id: int, message: dict):
         """Отправка сообщения конкретному сотруднику"""
-        logger.info(f"Broadcasting to employee {employee_id}: {message}")
-        logger.info(f"Number of employee connections: {len(self.active_connections['employee'])}")
-        
-        if not self.active_connections["employee"]:
-            logger.warning("No employee connections available")
-            return
-            
-        disconnected = []
-        
-        for connection in self.active_connections["employee"]:
+        websocket = self.employee_connections.get(employee_id)
+        if websocket:
             try:
-                await connection.send_json(message)
-                logger.info("Message sent successfully to employee")
+                await websocket.send_json(message)
             except Exception as e:
-                logger.error(f"Error sending message to employee: {e}")
-                disconnected.append(connection)
-                continue
-        
-        # Удаляем отключенные соединения
-        for connection in disconnected:
-            self.disconnect(connection, "employee")
+                logger.error(f"Error sending message to employee {employee_id}: {e}")
+                self.disconnect(websocket, "employee", employee_id)
 
     async def handle_ping(self, websocket: WebSocket):
-        """Обработка ping сообщений"""
+        """Обработка ping-сообщений"""
         try:
             await websocket.send_json({"type": "pong"})
-            logger.debug("Sent pong response")
         except Exception as e:
-            logger.error(f"Error sending pong: {e}")
+            logger.error(f"Error handling ping: {e}")
 
+    async def admin_endpoint(self, websocket: WebSocket):
+        """WebSocket endpoint для админов"""
+        await self.connect(websocket, "admin")
+        try:
+            while True:
+                data = await websocket.receive_json()
+                if data.get("type") == "ping":
+                    await self.handle_ping(websocket)
+        except Exception as e:
+            logger.error(f"Error in admin websocket: {e}")
+        finally:
+            self.disconnect(websocket, "admin")
+
+    async def employee_endpoint(self, websocket: WebSocket, employee_id: int):
+        """WebSocket endpoint для сотрудников"""
+        await self.connect(websocket, "employee", employee_id)
+        try:
+            while True:
+                data = await websocket.receive_json()
+                if data.get("type") == "ping":
+                    await self.handle_ping(websocket)
+        except Exception as e:
+            logger.error(f"Error in employee websocket: {e}")
+        finally:
+            self.disconnect(websocket, "employee", employee_id)
+
+
+# Создаем глобальный экземпляр менеджера уведомлений
 notification_manager = NotificationManager()