import asyncio import functools import json from pathlib import Path from typing import Any from aiohttp import web from app.features.core.utils import gen_random from app.library.log import get_logger from app.library.router import Route, RouteType, get_routes from app.library.Services import Services from app.library.Utils import load_modules from .config import Config from .encoder import Encoder from .Events import Event, EventBus, Events from .ItemDTO import Item LOG = get_logger() class WebSocketHub: def __init__(self, encoder: Encoder): self._encoder: Encoder = encoder self._clients: dict[str, web.WebSocketResponse] = {} def add(self, sid: str, ws: web.WebSocketResponse) -> None: self._clients[sid] = ws def remove(self, sid: str) -> None: self._clients.pop(sid, None) async def emit(self, event: str, data: Any, to: str | None = None) -> None: payload: str = self._encoder.encode({"event": event, "data": data}) if to: await self._send(to, payload) return for sid in list(self._clients.keys()): await self._send(sid, payload) async def disconnect(self, sid: str) -> None: ws: web.WebSocketResponse | None = self._clients.pop(sid, None) if ws and not ws.closed: await ws.close() async def disconnect_all(self) -> None: for sid in list(self._clients.keys()): await self.disconnect(sid) async def _send(self, sid: str, payload: str) -> None: ws: web.WebSocketResponse | None = self._clients.get(sid) if not ws or ws.closed: self._clients.pop(sid, None) return try: await ws.send_str(payload) except ConnectionResetError: self._clients.pop(sid, None) class HttpSocket: """ This class is used to handle WebSocket events. """ config: Config sio: WebSocketHub di_context: dict[str, Any] = {} def __init__( self, root_path: Path, encoder: Encoder | None = None, config: Config | None = None, sio: WebSocketHub | None = None, ): self.config = config or Config.get_instance() self._notify = EventBus.get_instance() encoder = encoder or Encoder() self.sio = sio or WebSocketHub(encoder=encoder) self.rootPath: Path = root_path async def event_handler(e: Event, _, **kwargs): payload = json.loads(encoder.encode(e)) await self.sio.emit(event=e.event, data=payload, **kwargs) services: Services = Services.get_instance() services.add_all( { k: v for k, v in { "config": self.config, "sio": self.sio, "encoder": encoder, "notify": self._notify, "root_path": self.rootPath, }.items() if not services.has(k) } ) self._notify.subscribe("frontend", event_handler, f"{HttpSocket.__name__}.emit") @staticmethod def ws_event(func): """ Decorator to mark a method as a socket event. """ @functools.wraps(func) async def wrapper(*args, **kwargs): return await func(*args, **kwargs) wrapper_with_event: Any = wrapper wrapper_with_event._ws_event = func.__name__ return wrapper async def on_shutdown(self, _: web.Application): LOG.debug("Shutting down socket server.") await self.sio.disconnect_all() LOG.debug("Socket server shutdown complete.") def attach(self, app: web.Application): app.on_shutdown.append(self.on_shutdown) base_path: str = self.config.base_path.rstrip("/") ws_path: str = f"{base_path}/ws" if base_path else "/ws" async def event_handler(data: Event, _): if not (data and data.data): return queue = Services.get_instance().get("queue") if queue is not None: await queue.add(item=Item.format(data.data)) self._notify.subscribe(Events.ADD_URL, event_handler, f"{HttpSocket.__name__}.add") load_modules(self.rootPath, self.rootPath / "routes" / "socket") socket_routes: dict[str, Route] = {route.path: route for route in get_routes(RouteType.SOCKET).values()} for route in socket_routes.values(): if self.config.debug: LOG.debug( "Registered socket route '%s' for %s %s.", route.name, route.method.value if isinstance(route.method, RouteType) else route.method, route.path, ) async def handle_message(sid: str, message: str) -> None: try: payload = json.loads(message) except json.JSONDecodeError: LOG.debug("Invalid websocket payload received.") return event = payload.get("event") if isinstance(payload, dict) else None if not event or not isinstance(event, str): LOG.debug("Missing websocket event name.") return if not (route := socket_routes.get(event)): LOG.debug("Unknown WebSocket event '%s'.", event, extra={"sid": sid, "event": event}) return data = payload.get("data") if isinstance(payload, dict) else None await Services.get_instance().handle_async(route.handler, sid=sid, data=data, event=event) async def handle_connect(sid: str) -> None: if not (route := socket_routes.get("connect")): return await Services.get_instance().handle_async(route.handler, sid=sid, event="connect") async def handle_disconnect(sid: str, reason: str | None) -> None: if not (route := socket_routes.get("disconnect")): return await Services.get_instance().handle_async(route.handler, sid=sid, data=reason, event="disconnect") async def ws_handler(request: web.Request) -> web.WebSocketResponse: ws = web.WebSocketResponse(heartbeat=10) await ws.prepare(request) sid: str = gen_random(14) self.sio.add(sid, ws) LOG.debug("WebSocket client '%s' connected.", sid, extra={"sid": sid}) await handle_connect(sid) async def handle_message_safe(sid: str, message: str) -> None: """Wrapper to handle message with error logging""" try: await handle_message(sid, message) except Exception as e: LOG.exception( "Failed to handle WebSocket message from client '%s'.", sid, extra={"sid": sid, "message_size": len(message), "exception_type": type(e).__name__}, ) try: i: int = 0 async for msg in ws: if msg.type == web.WSMsgType.TEXT: i = i + 1 asyncio.create_task(handle_message_safe(sid, msg.data), name=f"ws_msg_{sid}_{i}") elif msg.type == web.WSMsgType.ERROR: LOG.error( "WebSocket connection closed with exception %s", ws.exception(), extra={ "sid": sid, "exception": str(ws.exception()) if ws.exception() else None, }, ) finally: await handle_disconnect(sid, getattr(ws, "close_reason", None)) self.sio.remove(sid) LOG.debug("WebSocket client '%s' disconnected.", sid, extra={"sid": sid}) return ws if self.config.debug: LOG.debug("Registered WebSocket endpoint at '%s'.", ws_path, extra={"method": "GET", "path": ws_path}) app.router.add_get(ws_path, ws_handler, name="ws")