fix(wsrouter): ensure no message reply is lost (#906)

2 files changed, 16 insertions, 1 deletions

diff --git a/backend/decky_loader/plugin/messages.py b/backend/decky_loader/plugin/messages.py
index d53efbee..8e9bdf7a 100644
--- a/backend/decky_loader/plugin/messages.py
+++ b/backend/decky_loader/plugin/messages.py
@@ -19,18 +19,27 @@ class MethodCallResponse:
         self.success = success
         self.result = result
 
+class PluginStopped(Exception):
+    pass
+
 class MethodCallRequest:
     def __init__(self) -> None:
         self.id = str(uuid4())
         self.event = Event()
-        self.response: MethodCallResponse
+        self.response: MethodCallResponse | PluginStopped
     
     def set_result(self, dc: SocketResponseDict):
         self.response = MethodCallResponse(dc["success"], dc["res"])
         self.event.set()
+
+    def cancel(self):
+        self.response = PluginStopped("Plugin has been stopped")
+        self.event.set()
     
     async def wait_for_result(self):
         await self.event.wait()
+        if isinstance(self.response, PluginStopped):
+            raise self.response
         if not self.response.success:
             raise Exception(self.response.result)
         return self.response.result
\ No newline at end of file
diff --git a/backend/decky_loader/plugin/plugin.py b/backend/decky_loader/plugin/plugin.py
index a7edaa45..8648a5d6 100644
--- a/backend/decky_loader/plugin/plugin.py
+++ b/backend/decky_loader/plugin/plugin.py
@@ -94,6 +94,12 @@ class PluginWrapper:
             except CancelledError:
                 self.log.info(f"Stopping response listener for {self.name}")
                 await self._socket.close_socket_connection()
+
+                # Cancel the request so that WSRouter can perform cleanup
+                for request in self._method_call_requests.values():
+                    request.cancel()
+                
+                self._method_call_requests = {}
                 raise
             except:
                 pass