From 91abff7811c6db3a80c9426c0c9a0a247a3ada37 Mon Sep 17 00:00:00 2001 From: Lukas Senionis Date: Tue, 14 Jul 2026 05:02:13 +0300 Subject: fix(wsrouter): ensure no message reply is lost (#906) --- backend/decky_loader/plugin/messages.py | 11 +++- backend/decky_loader/plugin/plugin.py | 6 +++ backend/decky_loader/wsrouter.py | 95 ++++++++++++++++++++++----------- frontend/src/wsrouter.ts | 86 +++++++++++++++++++++-------- 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 { resolve: (res: T) => void; reject: (error: PyError) => void; promise: Promise; + message: CallMessage; } export class WSRouter extends Logger { runningCalls: Map> = new Map(); eventListeners: Map any>> = new Map(); ws?: WebSocket; - connectPromise?: Promise; // 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((resolve) => { + return new Promise((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(): PromiseResolver { - let resolver: Partial> = {}; + createPromiseResolver(message: CallMessage): PromiseResolver { + let resolver: Partial> = { message }; const promise = new Promise((resolve, reject) => { resolver.resolve = resolve; resolver.reject = reject; @@ -104,11 +128,6 @@ export class WSRouter extends Logger { return resolver as PromiseResolver; } - 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(route: string, ...args: Args): Promise { - const resolver = this.createPromiseResolver(); - const id = ++this.reqId; - + const resolver = this.createPromiseResolver({ 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)); + } + } } -- cgit v1.2.3