From 6253ec4aef664327880aab1f95d200a7488173ae Mon Sep 17 00:00:00 2001 From: comfyanonymous Date: Thu, 15 Jun 2023 11:01:06 -0400 Subject: [PATCH] Fix server crashing because of terminated websocket connection. --- server.py | 13 +++++++++---- 1 file changed, 9 insertions(+), 4 deletions(-) diff --git a/server.py b/server.py index 300221f6..f385cefb 100644 --- a/server.py +++ b/server.py @@ -30,6 +30,11 @@ import comfy.model_management class BinaryEventTypes: PREVIEW_IMAGE = 1 +async def send_socket_catch_exception(function, message): + try: + await function(message) + except (aiohttp.ClientError, aiohttp.ClientPayloadError, ConnectionResetError) as err: + print("send error:", err) @web.middleware async def cache_control(request: web.Request, handler): @@ -487,18 +492,18 @@ class PromptServer(): if sid is None: for ws in self.sockets.values(): - await ws.send_bytes(message) + await send_socket_catch_exception(ws.send_bytes, message) elif sid in self.sockets: - await self.sockets[sid].send_bytes(message) + await send_socket_catch_exception(self.sockets[sid].send_bytes, message) async def send_json(self, event, data, sid=None): message = {"type": event, "data": data} if sid is None: for ws in self.sockets.values(): - await ws.send_json(message) + await send_socket_catch_exception(ws.send_json, message) elif sid in self.sockets: - await self.sockets[sid].send_json(message) + await send_socket_catch_exception(self.sockets[sid].send_json, message) def send_sync(self, event, data, sid=None): self.loop.call_soon_threadsafe(