Show status errors when (temporarily) disconnected from the comfy server

- also improve error handling when managed server is stopped or crashes

ai_diffusion/comfy_client.py

@@ -234,7 +234,6 @@ async def _listen(self):
                 await self._listen_websocket(websocket)
             except websockets_exceptions.ConnectionClosedError as e:
                 log.warning(f"Websocket connection closed: {str(e)}")
-                await self._report(ClientEvent.disconnected, "")
             except OSError as e:
                 msg = _("Could not connect to websocket server at") + f"{url}: {str(e)}"
                 await self._report(ClientEvent.error, "", error=msg)
@@ -246,6 +245,8 @@ async def _listen(self):
             except Exception as e:
                 log.exception("Unhandled exception in websocket listener")
                 await self._report(ClientEvent.error, "", error=str(e))
+            finally:
+                await self._report(ClientEvent.disconnected, "")
 
     async def _listen_websocket(self, websocket: websockets_client.WebSocketClientProtocol):
         progress: Progress | None = None
@@ -358,7 +359,6 @@ async def disconnect(self):
             await asyncio.gather(
                 self._job_runner,
                 self._websocket_listener,
-                self._report(ClientEvent.disconnected, ""),
                 self._unsubscribe_workflows(),
             )
 

ai_diffusion/model.py

@@ -99,6 +99,7 @@ def __init__(self, document: Document, connection: Connection, workflows: Workfl
         self.jobs.selection_changed.connect(self.update_preview)
         self.error_changed.connect(lambda: self.has_error_changed.emit(self.has_error))
         connection.state_changed.connect(self._init_on_connect)
+        connection.error_changed.connect(self._forward_error)
         Styles.list().changed.connect(self._init_on_connect)
         self._init_on_connect()
 
@@ -110,6 +111,9 @@ def _init_on_connect(self):
             if self.upscale.upscaler == "":
                 self.upscale.upscaler = client.models.default_upscaler
 
+    def _forward_error(self, error: str):
+        self.error = error
+
     def generate(self):
         """Enqueue image generation for the current setup."""
         ok, msg = self._doc.check_color_mode()

ai_diffusion/resources.py

@@ -39,7 +39,7 @@ class CustomNode(NamedTuple):
         "External Tooling Nodes",
         "comfyui-tooling-nodes",
         "https://github.com/Acly/comfyui-tooling-nodes",
-        "9a9cbe78a5851a49da0b38bc9b17504837476e8f",
+        "0b01696f5b5da436b0f3e15660c13e9b2089ece6",
         ["ETN_LoadImageBase64", "ETN_LoadMaskBase64", "ETN_SendImageWebSocket", "ETN_Translate"],
     ),
     CustomNode(

ai_diffusion/server.py

@@ -438,8 +438,19 @@ async def run(self):
             async for line in self._process.stdout:
                 server_log.info(_decode_utf8_log_error(line).strip())
 
+            code = await asyncio.wait_for(self._process.wait(), timeout=1)
+            if code != 0:
+                log.error(f"Server process terminated with code {self._process.returncode}")
+            elif code is not None:
+                log.info(f"Server process was shut down sucessfully")
+
         except asyncio.CancelledError:
             pass
+        except asyncio.TimeoutError:
+            log.warning("Server process did not terminate after the pipe was closed")
+
+        self.state = ServerState.stopped
+        self._process = None
 
     async def stop(self):
         assert self.state is ServerState.running
@@ -448,15 +459,10 @@ async def stop(self):
                 log.info("Stopping server")
                 self._process.terminate()
                 self._task.cancel()
-                await asyncio.wait_for(self._process.communicate(), timeout=5)
-                log.info(f"Server terminated with code {self._process.returncode}")
+                await asyncio.wait_for(self._task, timeout=5)
         except asyncio.TimeoutError:
             log.warning("Server did not terminate in time")
             pass
-        finally:
-            self.state = ServerState.stopped
-            self._process = None
-            self._task = None
 
     def terminate(self):
         try:

ai_diffusion/ui/settings.py

@@ -304,7 +304,7 @@ def server_mode(self, mode: ServerMode):
         self._server_stack.setCurrentWidget(widget)
 
     def update_ui(self):
-        self._server_widget.update()
+        self._server_widget.update_ui()
 
     def _read(self):
         self.server_mode = settings.server_mode
@@ -854,7 +854,9 @@ def restore_defaults(self):
 
     def show(self, style: Optional[Style] = None):
         self.read()
+        self.connection.update_ui()
         super().show()
+
         if style:
             self._list.setCurrentRow(1)
             self.styles.current_style = style