diff options
Diffstat (limited to 'backend/decky_loader/wsrouter.py')
| -rw-r--r-- | backend/decky_loader/wsrouter.py | 47 |
1 files changed, 40 insertions, 7 deletions
diff --git a/backend/decky_loader/wsrouter.py b/backend/decky_loader/wsrouter.py index c501025e..84b83179 100644 --- a/backend/decky_loader/wsrouter.py +++ b/backend/decky_loader/wsrouter.py @@ -1,11 +1,32 @@ from logging import getLogger -from aiohttp import web +from asyncio import Future +from aiohttp import web, WSMsgType + +from enum import Enum + +from typing import Dict, Any, Callable + +from traceback import format_exc + +class MessageType(Enum): + # Call-reply + CALL = 0 + REPLY = 1 + ERROR = 2 + # Pub/sub + +running_calls: Dict[str, Future] = {} + +subscriptions: Dict[str, Callable[[Any]]] + +# {type: MessageType, data: dta, id: id} class WSRouter: def __init__(self) -> None: self.ws = None + self.req_id = 0 self.routes = {} self.logger = getLogger("WSRouter") @@ -23,16 +44,28 @@ class WSRouter: try: async for msg in ws: self.logger.debug(msg) - if msg.type == aiohttp.WSMsgType.TEXT: + if msg.type == WSMsgType.TEXT: self.logger.debug(msg.data) if msg.data == 'close': - # DO NOT RELY ON THIS! + # TODO DO NOT RELY ON THIS! break else: - # do stuff with the message - data = msg.json() - if self.routes[data.route]: - res = await self.routes[data.route](data.data) + match data.type: + case MessageType.CALL: + # do stuff with the message + data = msg.json() + if self.routes[data.route]: + try: + res = await self.routes[data.route](data.data) + await ws.send_json({type: MessageType.REPLY, id: data.id, data: res}) + except: + await ws.send_json({type: MessageType.ERROR, id: data.ud, data: format_exc()}) + case MessageType.REPLY: + if running_calls[data.id]: + running_calls[data.id].set_result(data.data) + case MessageType.ERROR: + if running_calls[data.id]: + running_calls[data.id].set_exception(data.data) finally: try: await ws.close() |
