diff --git a/projects/jupyter-server-ydoc/jupyter_server_ydoc/app.py b/projects/jupyter-server-ydoc/jupyter_server_ydoc/app.py index 8ee23ee3..70eaad44 100644 --- a/projects/jupyter-server-ydoc/jupyter_server_ydoc/app.py +++ b/projects/jupyter-server-ydoc/jupyter_server_ydoc/app.py @@ -10,7 +10,8 @@ from jupyter_ydoc.ybasedoc import YBaseDoc from pycrdt import Doc from pycrdt_websocket.ystore import BaseYStore -from traitlets import Bool, Float, Type +from pycrdt_websocket.websocket_server import exception_logger +from traitlets import Bool, Float, Type, Set, Instance from .handlers import DocSessionHandler, YDocWebSocketHandler from .loaders import FileLoaderMapping @@ -24,7 +25,6 @@ ) from .websocketserver import JupyterWebsocketServer, RoomNotFound - class YDocExtension(ExtensionApp): name = "jupyter_server_ydoc" app_name = "Collaboration" @@ -75,6 +75,10 @@ class YDocExtension(ExtensionApp): model.""", ) + _running_ywebsocket_server = Instance(asyncio.Task, allow_none=True) + ywebsocket_server = Instance(JupyterWebsocketServer, allow_none=True) + + def initialize(self): super().initialize() self.serverapp.event_logger.register_event_schema(EVENTS_SCHEMA_PATH) @@ -90,6 +94,14 @@ def initialize_settings(self): } ) + self.ywebsocket_server = JupyterWebsocketServer( + rooms_ready=False, + auto_clean_rooms=False, + ystore_class=self.ystore_class, + log=self.log, + exception_logger=exception_logger + ) + def initialize_handlers(self): self.serverapp.web_app.settings.setdefault( "page_config_data", @@ -103,13 +115,6 @@ def initialize_handlers(self): for k, v in self.config.get(self.ystore_class.__name__, {}).items(): setattr(self.ystore_class, k, v) - self.ywebsocket_server = JupyterWebsocketServer( - rooms_ready=False, - auto_clean_rooms=False, - ystore_class=self.ystore_class, - log=self.log, - ) - # self.settings is local to the ExtensionApp but here we need # the global app settings in which the file id manager will register # itself maybe at a later time. @@ -134,6 +139,14 @@ def initialize_handlers(self): ] ) + async def start_extension(self): + """Start the y-websocket server. + """ + self.log.info("Starting the Collaborative Document Server.") + + # Start the websocket server + self._running_ywebsocket_server = asyncio.create_task(self.ywebsocket_server.start()) + async def get_document( self: YDocExtension, *, diff --git a/projects/jupyter-server-ydoc/jupyter_server_ydoc/handlers.py b/projects/jupyter-server-ydoc/jupyter_server_ydoc/handlers.py index 59849058..9cdc9282 100644 --- a/projects/jupyter-server-ydoc/jupyter_server_ydoc/handlers.py +++ b/projects/jupyter-server-ydoc/jupyter_server_ydoc/handlers.py @@ -70,9 +70,7 @@ def create_task(self, aw): task.add_done_callback(self._background_tasks.discard) async def prepare(self): - if not self._websocket_server.started.is_set(): - self.create_task(self._websocket_server.start()) - await self._websocket_server.started.wait() + await self._websocket_server.started.wait() # Get room self._room_id: str = room_id_from_encoded_path(self.request.path)