initial WSRouter implementation

2 files changed, 50 insertions, 20 deletions

diff --git a/backend/decky_loader/loader.py b/backend/decky_loader/loader.py
index fa855c0a..ea3a2e75 100644
--- a/backend/decky_loader/loader.py
+++ b/backend/decky_loader/loader.py
@@ -87,9 +87,11 @@ class Loader:
             self.observer.schedule(self.watcher, self.plugin_path, recursive=True) # type: ignore
             self.observer.start()
             self.loop.create_task(self.enable_reload_wait())
-            
+
+        self.ws = WSRouter()
+
         server_instance.web_app.add_routes([
-            web.get("/socket", self.handle_ws),
+            web.get("/socket", self.ws.handle),
             web.get("/frontend/{path:.*}", self.handle_frontend_assets),
             web.get("/locales/{path:.*}", self.handle_frontend_locales),
             web.get("/plugins", self.get_plugins),
@@ -99,24 +101,6 @@ class Loader:
             web.post("/plugins/{plugin_name}/reload", self.handle_backend_reload_request)
         ])
 
-    async def handle_ws(self, request: web.Request):
-        self.logger.debug('Websocket connection starting')
-        ws = web.WebSocketResponse()
-        await ws.prepare(request)
-        self.logger.debug('Websocket connection ready')
-
-        async for msg in ws:
-            self.logger.debug(msg)
-            if msg.type == WSMsgType.TEXT:
-                self.logger.debug(msg.data)
-                if msg.data == 'close':
-                    await ws.close()
-                else:
-                    await ws.send_str(msg.data + '/answer')
-
-        self.logger.debug('Websocket connection closed')
-        return ws
-
     async def enable_reload_wait(self):
         if self.live_reload:
             await sleep(10)
diff --git a/backend/decky_loader/wsrouter.py b/backend/decky_loader/wsrouter.py
new file mode 100644
index 00000000..c501025e
--- /dev/null
+++ b/backend/decky_loader/wsrouter.py
@@ -0,0 +1,46 @@
+from logging import getLogger
+
+from aiohttp import web
+
+
+class WSRouter:
+    def __init__(self) -> None:
+        self.ws = None
+        self.routes = {}
+        self.logger = getLogger("WSRouter")
+
+    async def add_routes(self, routes):
+        self.routes.update(routes)
+
+    async def handle(self, request):
+        self.logger.debug('Websocket connection starting')
+        ws = web.WebSocketResponse()
+        await ws.prepare(request)
+        self.logger.debug('Websocket connection ready')
+
+        self.ws = ws
+        
+        try:
+            async for msg in ws:
+                self.logger.debug(msg)
+                if msg.type == aiohttp.WSMsgType.TEXT:
+                    self.logger.debug(msg.data)
+                    if msg.data == 'close':
+                        # 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)
+        finally:
+            try:
+                await ws.close()
+            except:
+                pass
+
+        self.logger.debug('Websocket connection closed')
+        return ws
+
+    async def write(self, dta: Dict[str, any]):
+        await self.ws.send_json(dta)
\ No newline at end of file