| 8 | |
| 9 | |
| 10 | class WSController(BaseController): |
| 11 | def __init__(self, app: FastAPI, container: Container) -> None: |
| 12 | super().__init__(container=container) |
| 13 | app.add_api_websocket_route(path="/ws/{user_id}", endpoint=self.wss) |
| 14 | |
| 15 | async def wss( |
| 16 | self, |
| 17 | websocket: WebSocket, |
| 18 | user_id: str, |
| 19 | ) -> None: |
| 20 | ws_service = self.container.ws_service() |
| 21 | try: |
| 22 | await ws_service.add_connection(user_id, websocket) |
| 23 | async for data in websocket.iter_text(): |
| 24 | try: |
| 25 | message = json.loads(data) |
| 26 | await ws_service.process_message(user_id=user_id, message=message, websocket=websocket) |
| 27 | except WebSocketDisconnect: |
| 28 | break |
| 29 | except json.JSONDecodeError: |
| 30 | self.logger.warning(f"Invalid JSON received from user {user_id}") |
| 31 | await ws_service.ws_manager.send_to_user(user_id=user_id, data={"message": "Incoming message is not valid"}, websocket=websocket) |
| 32 | except Exception as e: |
| 33 | self.logger.error(f"Error handling WebSocket message from user {user_id}: {e}") |
| 34 | break |
| 35 | |
| 36 | except Exception as e: |
| 37 | self.logger.error(f"WebSocket error for user {user_id}: {e}") |
| 38 | finally: |
| 39 | await ws_service.remove_connection(user_id=user_id, websocket=websocket) |