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

3 files changed, 80 insertions, 32 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
diff --git a/backend/decky_loader/wsrouter.py b/backend/decky_loader/wsrouter.py
index 691abd05..b0610b31 100644
--- a/backend/decky_loader/wsrouter.py
+++ b/backend/decky_loader/wsrouter.py
@@ -1,25 +1,28 @@
 from logging import getLogger
 
 from asyncio import AbstractEventLoop
-
 from aiohttp import WSCloseCode, WSMsgType, WSMessage
 from aiohttp.web import Application, WebSocketResponse, Request, Response, get
 
 from enum import IntEnum
-
 from typing import Callable, Coroutine, Dict, Any, cast
-
 from traceback import format_exc
 
 from .helpers import get_csrf_token
+from .plugin.messages import PluginStopped
 
 class MessageType(IntEnum):
     ERROR = -1
-    # Call-reply, Frontend -> Backend -> Frontend
+    # Call-reply, Frontend (CALL) -> Backend (REPLY|DISCARD) -> Frontend (RECEIVED_RESPONSE) -> Backend
     CALL = 0
     REPLY = 1
+    DISCARD = 2
+    RECEIVED_RESPONSE = 3
+    # Call-reply sync on connection lost,
+    #   Frontend (FULL_SYNC) -> Backend (REPLY|DISCARD...) -> Frontend (RECEIVED_RESPONSE...) -> Backend
+    FULL_SYNC = 4
     # Pub/Sub, Backend -> Frontend
-    EVENT = 3
+    EVENT = 5
 
 # WSMessage with slightly better typings
 class WSMessageExtra(WSMessage):
@@ -35,9 +38,8 @@ class WSRouter:
     def __init__(self, loop: AbstractEventLoop, server_instance: Application) -> None:
         self.loop = loop
         self.ws: WebSocketResponse | None = None
-        self.instance_id = 0
         self.routes: Dict[str, Route]  = {}
-        # self.subscriptions: Dict[str, Callable[[Any]]] = {}
+        self.pending_responses: Dict[int, Dict[str, Any] | None] = {}
         self.logger = getLogger("WSRouter")
 
         server_instance.add_routes([
@@ -45,9 +47,16 @@ class WSRouter:
         ])
 
     async def write(self, data: Dict[str, Any]):
+        # Cache all of the pending responses in case the connection is lost.
+        # Cleanup will happen during full-sync or when frontend confirms it has received the message
+        can_drop_message = True
+        if "id" in data and data["id"] in self.pending_responses:
+            self.pending_responses[data["id"]] = data
+            can_drop_message = False
+
         if self.ws != None:
             await self.ws.send_json(data)
-        else:
+        elif can_drop_message:
             self.logger.warning("Dropping message as there is no connected socket: %s", data)
 
     def add_route(self, name: str, route: Route):
@@ -57,26 +66,16 @@ class WSRouter:
         del self.routes[name]
 
     async def _call_route(self, route: str, args: ..., call_id: int):
-        instance_id = self.instance_id
-        error = None
         try:
             res = await self.routes[route](*args)
+            message = {"type": MessageType.REPLY.value, "id": call_id, "result": res}
+        except PluginStopped as err:
+            message = {"type": MessageType.DISCARD.value, "id": call_id}
         except Exception as err:
             error = {"name":err.__class__.__name__, "message":str(err), "traceback":format_exc()}
-            res = None
+            message = {"type": MessageType.ERROR.value, "id": call_id, "error": error}
         
-        if instance_id != self.instance_id:
-            try:
-                self.logger.warning("Ignoring %s reply from stale instance %d with args %s and response %s", route, instance_id, args, res)
-            except:
-                self.logger.warning("Ignoring %s reply from stale instance %d (failed to log event data)", route, instance_id)
-            finally:
-                return 
-
-        if error:
-            await self.write({"type": MessageType.ERROR.value, "id": call_id, "error": error})
-        else:
-            await self.write({"type": MessageType.REPLY.value, "id": call_id, "result": res})
+        await self.write(message)
 
     async def handle(self, request: Request):
         # Auth is a query param as JS WebSocket doesn't support headers
@@ -85,7 +84,6 @@ class WSRouter:
         self.logger.debug('Websocket connection starting')
         ws = WebSocketResponse()
         await ws.prepare(request)
-        self.instance_id += 1
         self.logger.debug('Websocket connection ready')
 
         if self.ws != None:
@@ -109,13 +107,11 @@ class WSRouter:
                         data = msg.json()
                         match data["type"]:
                             case MessageType.CALL.value:
-                                # do stuff with the message
-                                if data["route"] in self.routes:
-                                    self.logger.debug(f'Started PY call {data["route"]} ID {data["id"]}')
-                                    self.loop.create_task(self._call_route(data["route"], data["args"], data["id"]))
-                                else:
-                                    error = {"error":f'Route {data["route"]} does not exist.', "name": "RouteNotFoundError", "traceback": None}
-                                    self.loop.create_task(self.write({"type": MessageType.ERROR.value, "id": data["id"], "error": error}))
+                                self.handle_call_message(data)
+                            case MessageType.RECEIVED_RESPONSE.value:
+                                self.handle_received_response_message(data["id"])
+                            case MessageType.FULL_SYNC.value:
+                                self.handle_full_sync_message(data)
                             case _:
                                 self.logger.error("Unknown message type", data)
         finally:
@@ -128,6 +124,43 @@ class WSRouter:
         self.logger.debug('Websocket connection closed')
         return ws
 
+    def handle_call_message(self, data: Dict[str, Any]):
+        call_id = data["id"]
+        if call_id in self.pending_responses:
+            reply = self.pending_responses[call_id]
+            if reply:
+                self.logger.debug(f'Found pending PY call matching ID {call_id}. Sending reply...')
+                self.loop.create_task(self.write(reply))
+            else:
+                self.logger.debug(f'Found pending PY call matching ID {call_id}. Waiting for reply from plugin...')
+
+            return
+        
+        # Prepare a call entry for caching the reply once we have it, this will also
+        # help us identify (above) if we are already handling the call message or not. 
+        self.pending_responses[call_id] = None
+        if data["route"] in self.routes:
+            self.logger.debug(f'Started PY call {data["route"]} ID {call_id}')
+            self.loop.create_task(self._call_route(data["route"], data["args"], call_id))
+        else:
+            error = {"error":f'Route {data["route"]} does not exist.', "name": "RouteNotFoundError", "traceback": None}
+            self.loop.create_task(self.write({"type": MessageType.ERROR.value, "id": call_id, "error": error}))
+
+    def handle_received_response_message(self, call_id: int):
+        self.logger.debug(f'Removing pending response with ID {call_id}')
+        self.pending_responses.pop(call_id, None)
+            
+    def handle_full_sync_message(self, data: Dict[str, Any]):
+        messages = data["messages"]
+        outdated_response_ids = set(self.pending_responses.keys())
+
+        for message in messages:
+            outdated_response_ids.discard(message["id"])
+            self.handle_call_message(message)
+
+        for outdated_id in outdated_response_ids:
+            self.handle_received_response_message(outdated_id)
+
     async def emit(self, event: str, *args: Any):
         self.logger.debug(f'Firing frontend event {event} with args {args}')
         await self.write({ "type": MessageType.EVENT.value, "event": event, "args": args })