fix: a failing websocket.send would drop the connection silently

solara/server/flask.py

@@ -65,11 +65,17 @@ def close(self):
 
     def send_text(self, data: str) -> None:
         with self.lock:
-            self.ws.send(data)
+            try:
+                self.ws.send(data)
+            except simple_websocket.ws.ConnectionClosed:
+                raise websocket.WebSocketDisconnect()
 
     def send_bytes(self, data: bytes) -> None:
         with self.lock:
-            self.ws.send(data)
+            try:
+                self.ws.send(data)
+            except simple_websocket.ws.ConnectionClosed:
+                raise websocket.WebSocketDisconnect()
 
     async def receive(self):
         from anyio import to_thread

solara/server/kernel.py

@@ -215,8 +215,20 @@ def send_websockets(websockets: Set[websocket.WebsocketWrapper], binary_msg):
     for ws in list(websockets):
         try:
             ws.send(binary_msg)
-        except:  # noqa
-            # in case of any issue, we simply remove it from the list
+        except websocket.WebSocketDisconnect:
+            # ignore the exception, we tried to send while websocket closed
+            # just remove it from the websocket set
+            try:
+                # websocket can be modified by another thread
+                websockets.remove(ws)
+            except KeyError:
+                pass  # already removed
+        except Exception as e:  # noqa
+            logger.exception("Error sending message: %s, closing websocket", e)
+            try:
+                ws.close()
+            except Exception as e:  # noqa
+                logger.exception("Error closing websocket: %s", e)
             try:
                 # websocket can be modified by another thread
                 websockets.remove(ws)

solara/server/starlette.py

@@ -15,6 +15,7 @@
 import starlette.websockets
 import uvicorn.server
 import websockets.legacy.http
+import websockets.exceptions
 
 from solara.server.utils import path_is_child_of
 
@@ -121,9 +122,31 @@ async def process_messages_task(self):
             while len(self.to_send) > 0:
                 first = self.to_send.pop(0)
                 if isinstance(first, bytes):
-                    await self.ws.send_bytes(first)
+                    await self._send_bytes_exc(first)
                 else:
-                    await self.ws.send_text(first)
+                    await self._send_text_exc(first)
+
+    async def _send_bytes_exc(self, data: bytes):
+        # make sures we catch the starlette/websockets specific exception
+        # and re-raise it as a websocket.WebSocketDisconnect
+        try:
+            await self.ws.send_bytes(data)
+        except websockets.exceptions.ConnectionClosed as e:
+            raise websocket.WebSocketDisconnect() from e
+        except RuntimeError as e:
+            # starlette throws a RuntimeError once you call send after the connection is closed
+            raise websocket.WebSocketDisconnect() from e
+
+    async def _send_text_exc(self, data: str):
+        # make sures we catch the starlette/websockets specific exception
+        # and re-raise it as a websocket.WebSocketDisconnect
+        try:
+            await self.ws.send_text(data)
+        except websockets.exceptions.ConnectionClosed as e:
+            raise websocket.WebSocketDisconnect() from e
+        except RuntimeError as e:
+            # starlette throws a RuntimeError once you call send after the connection is closed
+            raise websocket.WebSocketDisconnect() from e
 
     def close(self):
         if self.portal is None:
@@ -133,25 +156,25 @@ def close(self):
 
     def send_text(self, data: str) -> None:
         if self.portal is None:
-            task = self.event_loop.create_task(self.ws.send_text(data))
+            task = self.event_loop.create_task(self._send_text_exc(data))
             self.tasks.add(task)
             task.add_done_callback(self.tasks.discard)
         else:
             if settings.main.experimental_performance:
                 self.to_send.append(data)
             else:
-                self.portal.call(self.ws.send_bytes, data)  # type: ignore
+                self.portal.call(self._send_bytes_exc, data)  # type: ignore
 
     def send_bytes(self, data: bytes) -> None:
         if self.portal is None:
-            task = self.event_loop.create_task(self.ws.send_bytes(data))
+            task = self.event_loop.create_task(self._send_bytes_exc(data))
             self.tasks.add(task)
             task.add_done_callback(self.tasks.discard)
         else:
             if settings.main.experimental_performance:
                 self.to_send.append(data)
             else:
-                self.portal.call(self.ws.send_bytes, data)  # type: ignore
+                self.portal.call(self._send_bytes_exc, data)  # type: ignore
 
     async def receive(self):
         if self.portal is None: