only configure the classes if they haven't been previously configured

Author: Zsailer

jupyter_server_documents/app.py

@@ -2,8 +2,6 @@
 from traitlets.config import Config
 from traitlets import Instance, Type
 
-from nextgen_kernels_api.services.kernels.client_manager import KernelClientManager
-
 from .handlers import RouteHandler, FileIDIndexHandler
 from .websockets import YRoomWebsocket
 from .rooms.yroom_manager import YRoomManager
@@ -49,8 +47,6 @@ class ServerDocsApp(ExtensionApp):
 
     yroom_manager = Instance(klass=YRoomManager, allow_none=True)
 
-    client_manager = Instance(klass=KernelClientManager, allow_none=True)
-
     def initialize(self):
         super().initialize()
 
@@ -75,17 +71,6 @@ def get_fileid_manager():
         self.outputs_manager = self.outputs_manager_class(parent=self)
         self.settings["outputs_manager"] = self.outputs_manager
 
-        # Initialize KernelClientRegistry as singleton
-        # The KernelClientWebsocketConnection from nextgen-kernels-api will access this via .instance()
-        self.client_manager = KernelClientManager.instance(
-            parent=self,
-            multi_kernel_manager=self.serverapp.kernel_manager
-        )
-        self.settings["client_manager"] = self.client_manager
-
-        # Register event listener for client management
-        self.client_manager.register_event_listener(self.serverapp.event_logger)
-
         # Serve Jupyter Collaboration API on
         # `self.settings["jupyter_server_ydoc"]` for compatibility with
         # extensions depending on Jupyter Collaboration
@@ -95,35 +80,42 @@ def get_fileid_manager():
         )
 
     def _link_jupyter_server_extension(self, server_app):
-        """Setup custom config needed by this extension."""
+        """Setup custom config needed by this extension.
+
+        Only applies configuration if not already set by user config.
+        """
         c = Config()
-        # Use nextgen-kernels-api's multi-kernel manager
-        # This manager disables activity watching and buffering (which we don't need)
-        c.ServerApp.kernel_manager_class = "nextgen_kernels_api.services.kernels.kernelmanager.MultiKernelManager"
-        c.ServerApp.kernel_websocket_connection_class = 'nextgen_kernels_api.services.kernels.connection.kernel_client_connection.KernelClientWebsocketConnection'
-
-        # Configure MultiKernelManager to use nextgen's KernelManager
-        c.MultiKernelManager.kernel_manager_class = "nextgen_kernels_api.services.kernels.kernelmanager.KernelManager"
-
-        # Configure the KernelManager to use DocumentAwareKernelClient
-        c.KernelManager.client_class = "jupyter_server_documents.kernel_client.DocumentAwareKernelClient"
-        c.KernelManager.client_factory = "jupyter_server_documents.kernel_client.DocumentAwareKernelClient"
-
-        # Use custom session manager for YRoom integration
-        c.ServerApp.session_manager_class = "jupyter_server_documents.session_manager.YDocSessionManager"
-
-        # Configure websocket connection to exclude output messages
-        # Output messages are handled by DocumentAwareKernelClient's output processor
-        # and should not be sent to websocket clients to avoid duplicate processing
-        c.KernelClientWebsocketConnection.exclude_msg_types = [
-            ("status", "iopub"),
-            ("stream", "iopub"),
-            ("display_data", "iopub"),
-            ("execute_result", "iopub"),
-            ("error", "iopub"),
-            ("update_display_data", "iopub"),
-            ("clear_output", "iopub")
-        ]
+
+        # Configure kernel manager classes to use nextgen-kernels-api
+        if not server_app.config.ServerApp.get("kernel_manager_class"):
+            c.ServerApp.kernel_manager_class = "nextgen_kernels_api.services.kernels.kernelmanager.MultiKernelManager"
+
+        if not server_app.config.ServerApp.get("kernel_websocket_connection_class"):
+            c.ServerApp.kernel_websocket_connection_class = "nextgen_kernels_api.services.kernels.connection.kernel_client_connection.KernelClientWebsocketConnection"
+
+        if not server_app.config.ServerApp.get("session_manager_class"):
+            c.ServerApp.session_manager_class = "jupyter_server_documents.session_manager.YDocSessionManager"
+
+        # Configure kernel manager hierarchy
+        if not server_app.config.MultiKernelManager.get("kernel_manager_class"):
+            c.MultiKernelManager.kernel_manager_class = "nextgen_kernels_api.services.kernels.kernelmanager.KernelManager"
+
+        # Configure kernel client
+        if not server_app.config.KernelManager.get("client_class"):
+            c.KernelManager.client_class = "jupyter_server_documents.kernel_client.DocumentAwareKernelClient"
+            c.KernelManager.client_factory = "jupyter_server_documents.kernel_client.DocumentAwareKernelClient"
+
+        # Configure websocket message filtering
+        if not server_app.config.KernelClientWebsocketConnection.get("exclude_msg_types"):
+            c.KernelClientWebsocketConnection.exclude_msg_types = [
+                ("status", "iopub"),
+                ("stream", "iopub"),
+                ("display_data", "iopub"),
+                ("execute_result", "iopub"),
+                ("error", "iopub"),
+                ("update_display_data", "iopub"),
+                ("clear_output", "iopub"),
+            ]
 
         server_app.update_config(c)
         super()._link_jupyter_server_extension(server_app)

jupyter_server_documents/rooms/yroom_file_api.py

@@ -546,7 +546,7 @@ async def save(self, jupyter_ydoc: YBaseDoc):
 
             # Set most recent `last_modified` timestamp
             if file_data['last_modified']:
-                self.log.info(f"Reseting last_modified to {file_data['last_modified']}")
+                self.log.debug(f"Resetting last_modified to {file_data['last_modified']}")
                 self._last_modified = file_data['last_modified']
 
             # Set `dirty` to `False` to hide the "unsaved changes" icon in the

jupyter_server_documents/session_manager.py

@@ -6,7 +6,6 @@
 from jupyter_server_documents.rooms.yroom_manager import YRoomManager
 from jupyter_server_documents.rooms.yroom import YRoom
 from jupyter_server_documents.kernel_client import DocumentAwareKernelClient
-from nextgen_kernels_api.services.kernels.client_manager import KernelClientManager
 
 
 class YDocSessionManager(SessionManager): 
@@ -31,11 +30,6 @@ def yroom_manager(self) -> YRoomManager:
         """The Jupyter Server's YRoom Manager."""
         return self.serverapp.web_app.settings["yroom_manager"]
 
-    @property
-    def client_manager(self) -> KernelClientManager:
-        """The Kernel Client Manager."""
-        return self.serverapp.web_app.settings["client_manager"]
-
     _room_ids: dict[str, str]
     """
     Dictionary of room IDs, keyed by session ID.
@@ -83,7 +77,34 @@ async def create_session(
     ) -> dict[str, Any]:
         """
         After creating a session, connects the yroom to the kernel client.
+        Sets kernel status to "starting" before kernel launch.
         """
+        # For notebooks, set up the YRoom and set initial status before starting kernel
+        should_setup_yroom = (
+            type == "notebook" and
+            name is not None and
+            path is not None
+        )
+
+        yroom = None
+        if should_setup_yroom:
+            # Calculate the real path
+            real_path = os.path.join(os.path.split(path)[0], name)
+            
+            # Initialize the YRoom before starting the kernel
+            file_id = self.file_id_manager.index(real_path)
+            room_id = f"json:notebook:{file_id}"
+            yroom = self.yroom_manager.get_room(room_id)
+
+            # Set initial kernel status to "starting" in awareness
+            awareness = yroom.get_awareness()
+            if awareness is not None:
+                self.log.info("Setting kernel execution_state to 'starting' before kernel launch")
+                awareness.set_local_state_field(
+                    "kernel", {"execution_state": "starting"}
+                )
+
+        # Now create the session and start the kernel
         session_model = await super().create_session(
             path,
             name,
@@ -108,27 +129,21 @@ async def create_session(
             self.log.warning(f"`name` or `path` was not given for new session at '{path}'.")
             return session_model
 
-        # Otherwise, get a `YRoom` and add it to this session's kernel client.
-
-        # When JupyterLab creates a session, it uses a fake path
-        # which is the relative path + UUID, i.e. the notebook
-        # name is incorrect temporarily. It later makes multiple
-        # updates to the session to correct the path.
-        # 
-        # Here, we create the true path to store in the fileID service
-        # by dropping the UUID and appending the file name.
-        real_path = os.path.join(os.path.split(path)[0], name)
+        # Otherwise, add the YRoom to this session's kernel client.
 
-        # Get YRoom for this session and store its ID in `self._room_ids`
-        yroom = self._init_session_yroom(session_id, real_path)
+        # Store the room ID for this session
+        if yroom:
+            self._room_ids[session_id] = yroom.room_id
+        else:
+            # Shouldn't happen, but handle it anyway
+            real_path = os.path.join(os.path.split(path)[0], name)
+            yroom = self._init_session_yroom(session_id, real_path)
 
         # Add YRoom to this session's kernel client
-        # TODO: we likely have a race condition here... need to 
-        # think about it more. Currently, the kernel client gets
-        # created after the kernel starts fully. We need the 
-        # kernel client instantiated _before_ trying to connect
-        # the yroom.
-        kernel_client = self.client_manager.get_client(kernel_id)
+        # Ensure the kernel client is fully connected before proceeding
+        # to avoid queuing messages on first execution
+        kernel_manager = self.serverapp.kernel_manager.get_kernel(kernel_id)
+        kernel_client = kernel_manager.kernel_client
         await kernel_client.add_yroom(yroom)
         self.log.info(f"Connected yroom {yroom.room_id} to kernel {kernel_id}. yroom: {yroom}")
         return session_model
@@ -158,10 +173,12 @@ async def update_session(self, session_id: str, **update) -> None:
         )
         yroom = self.get_yroom(session_id)
         if old_kernel_id:
-            old_kernel_client = self.client_manager.get_client(old_kernel_id)
+            old_kernel_manager = self.serverapp.kernel_manager.get_kernel(old_kernel_id)
+            old_kernel_client = old_kernel_manager.kernel_client
             await old_kernel_client.remove_yroom(yroom=yroom)
         if new_kernel_id:
-            new_kernel_client = self.client_manager.get_client(new_kernel_id)
+            new_kernel_manager = self.serverapp.kernel_manager.get_kernel(new_kernel_id)
+            new_kernel_client = new_kernel_manager.kernel_client
             await new_kernel_client.add_yroom(yroom=yroom)
 
         # Apply update and return
@@ -177,7 +194,8 @@ async def delete_session(self, session_id):
 
         # Remove YRoom from session's kernel client
         yroom = self.get_yroom(session_id)
-        kernel_client = self.client_manager.get_client(kernel_id)
+        kernel_manager = self.serverapp.kernel_manager.get_kernel(kernel_id)
+        kernel_client = kernel_manager.kernel_client
         await kernel_client.remove_yroom(yroom)
 
         # Remove room ID stored for the session

pyproject.toml

@@ -32,7 +32,7 @@ dependencies = [
     "jupyter_server_fileid>=0.9.0,<0.10.0",
     "pycrdt>=0.12.0,<0.13.0",
     "jupyter_ydoc>=3.0.0,<4.0.0",
-    "nextgen-kernels-api>=0.6.0",
+    "nextgen-kernels-api>=0.9.0",
 ]
 dynamic = ["version", "description", "authors", "urls", "keywords"]
 
@@ -97,3 +97,7 @@ before-build-python = ["jlpm clean:all"]
 
 [tool.check-wheel-contents]
 ignore = ["W002"]
+
+[tool.pytest.ini_options]
+testpaths = ["jupyter_server_documents/tests"]
+norecursedirs = ["repos", ".git", ".pixi", "node_modules"]