Log (instead of raise) exceptions when running as a server extension

Author: Zsailer

projects/jupyter-server-ydoc/jupyter_server_ydoc/app.py

@@ -22,7 +22,7 @@
     encode_file_path,
     room_id_from_encoded_path,
 )
-from .websocketserver import JupyterWebsocketServer, RoomNotFound
+from .websocketserver import JupyterWebsocketServer, RoomNotFound, exception_logger
 
 
 class YDocExtension(ExtensionApp):
@@ -107,6 +107,9 @@ def initialize_handlers(self):
             rooms_ready=False,
             auto_clean_rooms=False,
             ystore_class=self.ystore_class,
+            # Log exceptions, because we don't want the websocket server
+            # to _ever_ crash permanently in a live jupyter_server.
+            exception_handler=exception_logger,
             log=self.log,
         )
 

projects/jupyter-server-ydoc/jupyter_server_ydoc/handlers.py

@@ -104,6 +104,23 @@ async def prepare(self):
                         self.log,
                         self._document_save_delay,
                     )
+                    
+                    def exception_logger(exception: Exception, log) -> bool:
+                        """A function that catches any exceptions raised in the websocket
+                        server and logs them. 
+
+                        The protects the websocket server's task group from cancelling
+                        anytime an exception is raised.
+                        """
+                        room_id = "unknown"
+                        if self.room.room_id:
+                            room_id = self.room.room_id
+                        log.error(f"Document Room Exception, (room_id={room_id}: ", exc_info=exception)
+                        return True
+
+                    # Logging exceptions, instead of raising them here to ensure
+                    # that the y-rooms stay alive even after an exception is seen.
+                    self.room.exception_handler = exception_logger
 
                 else:
                     # TransientRoom

projects/jupyter-server-ydoc/jupyter_server_ydoc/rooms.py

@@ -52,6 +52,7 @@ def __init__(
         self._document.observe(self._on_document_change)
         self._file.observe(self.room_id, self._on_outofband_change)
 
+
     @property
     def room_id(self) -> str:
         """

projects/jupyter-server-ydoc/jupyter_server_ydoc/websocketserver.py

@@ -5,7 +5,7 @@
 
 import asyncio
 from logging import Logger
-from typing import Any
+from typing import Any, Callable
 
 from pycrdt_websocket.websocket_server import WebsocketServer, YRoom
 from pycrdt_websocket.ystore import BaseYStore
@@ -16,6 +16,17 @@ class RoomNotFound(LookupError):
     pass
 
 
+def exception_logger(exception: Exception, log: Logger) -> bool:
+    """A function that catches any exceptions raised in the websocket
+    server and logs them. 
+
+    The protects the websocket server's task group from cancelling
+    anytime an exception is raised.
+    """
+    log.error("Jupyter Websocket Server: ", exc_info=exception)
+    return True
+
+
 class JupyterWebsocketServer(WebsocketServer):
     """Ypy websocket server.
 
@@ -30,9 +41,15 @@ def __init__(
         ystore_class: BaseYStore,
         rooms_ready: bool = True,
         auto_clean_rooms: bool = True,
+        exception_handler:  Callable[[Exception, Logger], bool] | None = None,
         log: Logger | None = None,
     ):
-        super().__init__(rooms_ready, auto_clean_rooms, log)
+        super().__init__(
+            rooms_ready=rooms_ready, 
+            auto_clean_rooms=auto_clean_rooms, 
+            exception_handler=exception_handler,
+            log=log, 
+        )
         self.ystore_class = ystore_class
         self.ypatch_nb = 0
         self.connected_users: dict[Any, Any] = {}

projects/jupyter-server-ydoc/pyproject.toml

@@ -30,7 +30,7 @@ authors = [
 dependencies = [
     "jupyter_server>=2.4.0,<3.0.0",
     "jupyter_ydoc>=2.0.0,<3.0.0",
-    "pycrdt-websocket>=0.13.0,<0.14.0",
+    "pycrdt-websocket>=0.13.1,<0.14.0",
     "jupyter_events>=0.10.0",
     "jupyter_server_fileid>=0.7.0,<1",
     "jsonschema>=4.18.0"