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

4 files changed, 144 insertions, 54 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 })
diff --git a/frontend/src/wsrouter.ts b/frontend/src/wsrouter.ts
index 00d0bd2e..5d02fc0c 100644
--- a/frontend/src/wsrouter.ts
+++ b/frontend/src/wsrouter.ts
@@ -8,11 +8,16 @@ declare global {
 
 enum MessageType {
   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,
 }
 
 interface CallMessage {
@@ -28,6 +33,21 @@ interface ReplyMessage {
   id: number;
 }
 
+interface DiscardMessage {
+  type: MessageType.DISCARD;
+  id: number;
+}
+
+interface ReceivedResponseMessage {
+  type: MessageType.RECEIVED_RESPONSE;
+  id: number;
+}
+
+interface FullSyncMessage {
+  type: MessageType.FULL_SYNC;
+  messages: CallMessage[];
+}
+
 interface ErrorMessage {
   type: MessageType.ERROR;
   error: { name: string; error: string; traceback: string | null };
@@ -58,20 +78,21 @@ interface EventMessage {
   args: any;
 }
 
-type Message = CallMessage | ReplyMessage | ErrorMessage | EventMessage;
+type MessageToBackend = CallMessage | ReceivedResponseMessage | FullSyncMessage;
+type MessageFromBackend = ReplyMessage | ErrorMessage | DiscardMessage | EventMessage;
 
 // Helper to resolve a promise from the outside
 interface PromiseResolver<T> {
   resolve: (res: T) => void;
   reject: (error: PyError) => void;
   promise: Promise<T>;
+  message: CallMessage;
 }
 
 export class WSRouter extends Logger {
   runningCalls: Map<number, PromiseResolver<any>> = new Map();
   eventListeners: Map<string, Set<(...args: any) => any>> = new Map();
   ws?: WebSocket;
-  connectPromise?: Promise<void>;
   // Used to map results and errors to calls
   reqId: number = 0;
   constructor() {
@@ -79,23 +100,26 @@ export class WSRouter extends Logger {
   }
 
   connect() {
-    return (this.connectPromise = new Promise<void>((resolve) => {
+    return new Promise<void>((resolve) => {
       // Auth is a query param as JS WebSocket doesn't support headers
       this.ws = new WebSocket(`ws://127.0.0.1:1337/ws?auth=${deckyAuthToken}`);
 
       this.ws.addEventListener('open', () => {
         this.debug('WS Connected');
         resolve();
-        delete this.connectPromise;
+
+        // Synchronize frontend and backend by forwarding all the running calls again
+        // and letting backend reply to them accordingly.
+        const messages = Array.from(this.runningCalls.values()).map((resolver) => resolver.message);
+        this.write({ type: MessageType.FULL_SYNC, messages });
       });
       this.ws.addEventListener('message', this.onMessage.bind(this));
       this.ws.addEventListener('close', this.onError.bind(this));
-      // this.ws.addEventListener('error', this.onError.bind(this));
-    }));
+    });
   }
 
-  createPromiseResolver<T>(): PromiseResolver<T> {
-    let resolver: Partial<PromiseResolver<T>> = {};
+  createPromiseResolver<T>(message: CallMessage): PromiseResolver<T> {
+    let resolver: Partial<PromiseResolver<T>> = { message };
     const promise = new Promise<T>((resolve, reject) => {
       resolver.resolve = resolve;
       resolver.reject = reject;
@@ -104,11 +128,6 @@ export class WSRouter extends Logger {
     return resolver as PromiseResolver<T>;
   }
 
-  async write(data: Message) {
-    if (this.connectPromise) await this.connectPromise;
-    this.ws?.send(JSON.stringify(data));
-  }
-
   addEventListener(event: string, listener: (...args: any) => any) {
     if (!this.eventListeners.has(event)) {
       this.eventListeners.set(event, new Set([listener]));
@@ -130,7 +149,7 @@ export class WSRouter extends Logger {
 
   async onMessage(msg: MessageEvent) {
     try {
-      const data = JSON.parse(msg.data) as Message;
+      const data = JSON.parse(msg.data) as MessageFromBackend;
       switch (data.type) {
         case MessageType.REPLY:
           if (this.runningCalls.has(data.id)) {
@@ -138,6 +157,8 @@ export class WSRouter extends Logger {
             this.runningCalls.delete(data.id);
             this.debug(`[${data.id}] Resolved PY call with value`, data.result);
           }
+
+          this.write({ type: MessageType.RECEIVED_RESPONSE, id: data.id });
           break;
 
         case MessageType.ERROR:
@@ -147,6 +168,21 @@ export class WSRouter extends Logger {
             this.runningCalls.delete(data.id);
             this.debug(`[${data.id}] Rejected PY call with error`, data.error);
           }
+
+          this.write({ type: MessageType.RECEIVED_RESPONSE, id: data.id });
+          break;
+
+        case MessageType.DISCARD:
+          if (this.runningCalls.has(data.id)) {
+            // We are explicitly not resolving the promise here as this message
+            // is only received (at the moment) when plugin is being unloaded and
+            // the promise should be garbage-collected, unless the plugin is doing
+            // very bad things!
+            this.runningCalls.delete(data.id);
+            this.debug(`[${data.id}] Discarding PY call`);
+          }
+
+          this.write({ type: MessageType.RECEIVED_RESPONSE, id: data.id });
           break;
 
         case MessageType.EVENT:
@@ -177,15 +213,12 @@ export class WSRouter extends Logger {
 
   // this.call<[number, number], string>('methodName', 1, 2);
   call<Args extends any[] = [], Return = void>(route: string, ...args: Args): Promise<Return> {
-    const resolver = this.createPromiseResolver<Return>();
-
     const id = ++this.reqId;
-
+    const resolver = this.createPromiseResolver<Return>({ type: MessageType.CALL, route, args, id });
     this.runningCalls.set(id, resolver);
 
     this.debug(`[${id}] Calling PY method ${route} with args`, args);
-
-    this.write({ type: MessageType.CALL, route, args, id });
+    this.write(resolver.message);
 
     return resolver.promise;
   }
@@ -196,8 +229,17 @@ export class WSRouter extends Logger {
 
   async onError(error: any) {
     this.error('WS DISCONNECTED', error);
-    // TODO queue up lost messages and send them once we connect again
     await sleep(5000);
     await this.connect();
   }
+
+  private write(data: MessageToBackend) {
+    // In case the message is discarded here:
+    //   - CallMessage will be resent during FULL_SYNC
+    //   - Backend cleanup via missed ReceivedReplyMessage will be done during FULL_SYNC
+    //   - FullSyncMessage will be resent during FULL_SYNC
+    if (this.ws?.readyState === WebSocket.OPEN) {
+      this.ws?.send(JSON.stringify(data));
+    }
+  }
 }