Reinject loader if steam got restarted

2 files changed, 44 insertions, 7 deletions

diff --git a/plugin_loader/injector.py b/plugin_loader/injector.py
index e130d830..6ca9e8e9 100644
--- a/plugin_loader/injector.py
+++ b/plugin_loader/injector.py
@@ -1,11 +1,13 @@
 #Injector code from https://github.com/SteamDeckHomebrew/steamdeck-ui-inject. More info on how it works there.
 
 from aiohttp import ClientSession
-from logging import debug
+from logging import debug, getLogger
 from asyncio import sleep
 
 BASE_ADDRESS = "http://localhost:8080"
 
+logger = getLogger("Injector")
+
 class Tab:
     def __init__(self, res) -> None:
         self.title = res["title"]
@@ -68,7 +70,7 @@ async def get_tabs():
                 res = await web.get("{}/json".format(BASE_ADDRESS))
                 break
             except:
-                print("Steam isn't available yet. Wait for a moment...")
+                logger.info("Steam isn't available yet. Wait for a moment...")
                 await sleep(5)
 
         if res.status == 200:
@@ -82,4 +84,16 @@ async def inject_to_tab(tab_name, js):
     tab = next((i for i in tabs if i.title == tab_name), None)
     if not tab:
         raise ValueError("Tab {} not found in running tabs".format(tab_name))
-    debug(await tab.evaluate_js(js))
+    logger.debug(f"Injected JavaScript Result: {await tab.evaluate_js(js)}")
+
+async def tab_has_element(tab_name, element_name):
+    tabs = await get_tabs()
+    tab = next((i for i in tabs if i.title == tab_name), None)
+    if not tab:
+        return False
+    res = await tab.evaluate_js(f"document.getElementById('{element_name}') != null")
+    
+    if not "result" in res or not "result" in res["result"] or not "value" in res["result"]["result"]:
+        return False;
+
+    return res["result"]["result"]["value"]
diff --git a/plugin_loader/main.py b/plugin_loader/main.py
index 7e13a2d8..a5d044b1 100644
--- a/plugin_loader/main.py
+++ b/plugin_loader/main.py
@@ -1,5 +1,6 @@
-from logging import basicConfig, INFO, DEBUG
+from logging import getLogger, basicConfig, INFO, DEBUG
 from os import getenv
+
 CONFIG = {
     "plugin_path": getenv("PLUGIN_PATH", "/home/deck/homebrew/plugins"),
     "server_host": getenv("SERVER_HOST", "127.0.0.1"),
@@ -13,13 +14,16 @@ from aiohttp.web import Application, run_app, static
 from aiohttp_jinja2 import setup as jinja_setup
 from jinja2 import FileSystemLoader
 from os import path
-from asyncio import get_event_loop
+from asyncio import get_event_loop, sleep
 from json import loads, dumps
 
 from loader import Loader
-from injector import inject_to_tab, get_tabs
+from injector import inject_to_tab, get_tabs, tab_has_element
 from utilities import util_methods
 
+
+logger = getLogger("Main")
+
 class PluginManager:
     def __init__(self) -> None:
         self.loop = get_event_loop()
@@ -30,6 +34,14 @@ class PluginManager:
         self.web_app.on_startup.append(self.inject_javascript)
         self.web_app.add_routes([static("/static", path.join(path.dirname(__file__), 'static'))])
         self.loop.create_task(self.method_call_listener())
+        self.loop.create_task(self.loader_reinjector())
+
+        self.loop.set_exception_handler(self.exception_handler)
+
+    def exception_handler(self, loop, context):
+        if context["message"] == "Unclosed connection":
+            return
+        loop.default_exception_handler(context)
 
     async def resolve_method_call(self, tab, call_id, response):
         await tab._send_devtools_cmd({
@@ -72,8 +84,19 @@ class PluginManager:
                 method = loads(data["params"]["args"][0]["value"])
                 self.loop.create_task(self.handle_method_call(method, tab))
                 
+    async def loader_reinjector(self):
+        while True:
+            await sleep(1)
+            if not await tab_has_element("QuickAccess", "plugin_iframe"):
+                logger.info("Plugin loader isn't present in Steam anymore, reinjecting...")
+                await self.inject_javascript()
+
     async def inject_javascript(self, request=None):
-        await inject_to_tab("QuickAccess", open(path.join(path.dirname(__file__), "static/plugin_page.js"), "r").read())
+        try:
+            await inject_to_tab("QuickAccess", open(path.join(path.dirname(__file__), "static/plugin_page.js"), "r").read())
+        except:
+            logger.info("Failed to inject JavaScript into tab")
+            pass
 
     def run(self):
         return run_app(self.web_app, host=CONFIG["server_host"], port=CONFIG["server_port"], loop=self.loop)