update kernel client to always use latest rooms from YRoomManager

Author: dlqqq

jupyter_server_documents/kernels/kernel_client.py

@@ -1,9 +1,9 @@
 """
 A new Kernel client that is aware of ydocuments.
 """
+from __future__ import annotations
 import anyio
 import asyncio
-import json
 import typing as t
 
 from traitlets import Set, Instance, Any, Type, default
@@ -16,6 +16,9 @@
 
 from .kernel_client_abc import AbstractDocumentAwareKernelClient
 
+if t.TYPE_CHECKING:
+    from jupyter_server_documents.rooms.yroom_manager import YRoomManager
+
 
 class DocumentAwareKernelClient(AsyncKernelClient):
     """
@@ -39,11 +42,6 @@ def _default_message_cache(self):
     # message is received.
     _listeners = Set(allow_none=True)
 
-    # A set of YRooms that will intercept output and kernel
-    # status messages.
-    _yrooms: t.Set[YRoom] = Set(trait=Instance(YRoom), default_value=set())
-
-
     output_processor = Instance(
         OutputProcessor,
         allow_none=True
@@ -58,6 +56,24 @@ def _default_message_cache(self):
     def _default_output_processor(self) -> OutputProcessor:
         self.log.info("Creating output processor")
         return self.output_process_class(parent=self, config=self.config)
+    
+    _yroom_manager: YRoomManager | None
+    """
+    The YRoomManager registered via `self.bind_yroom_manager()`, which must be
+    called before adding any `YRoom` via `self.add_yroom()`.
+    """
+
+    _yroom_ids: set[str]
+    """
+    The set of room IDs that are registered with this kernel client. This class
+    stores room IDs instead of `YRoom` instances because `YRoom` instances may
+    be deleted once inactive.
+    """
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self._yroom_manager = None
+        self._yroom_ids = set()
 
     async def start_listening(self):
         """Start listening to messages coming from the kernel.
@@ -286,18 +302,63 @@ async def handle_document_related_message(self, msg: t.List[bytes]) -> t.Optiona
         # Default return if message is processed and does not need forwarding
         return msg
 
+    @property
+    def _yrooms(self) -> list[YRoom]:
+        """
+        Returns the list of YRoom instances registered to this kernel client.
+        """
+        if len(self._yroom_ids) == 0:
+            return []
+
+        assert self._yroom_manager
+        rooms: list[YRoom] = []
+
+        # Always call `get_room()` to get the latest reference to the room. We
+        # must do this since rooms may be deleted upon inactivity. The
+        # `get_room()` method returns a cached value as long as the room was not
+        # deleted, so this is very fast in most cases.
+        for room_id in self._yroom_ids:
+            room = self._yroom_manager.get_room(room_id)
+            rooms.append(room)
+
+        return rooms
+
     async def add_yroom(self, yroom: YRoom):
         """
-        Register a YRoom with this kernel client. YRooms will
-        intercept display and kernel status messages.
+        Register a `YRoom` with this kernel client, given the room ID.
+        Registered `YRoom`s will intercept display and kernel status messages.
+
+        `self.bind_yroom_manager()` must be called before using this method.
         """
-        self._yrooms.add(yroom)
+        assert self._yroom_manager
+        self._yroom_ids.add(yroom.room_id)
+        self.log.info(
+            f"Added room '{yroom.room_id}' to kernel '{self.kernel_name}'. "
+            f"Total rooms: {len(self._yroom_ids)}"
+        )
 
     async def remove_yroom(self, yroom: YRoom):
         """
-        De-register a YRoom from handling kernel client messages.
+        De-register a `YRoom` from this kernel client, given the room ID.
+
+        `self.bind_yroom_manager()` must be called before using this method.
+        """
+        self._yrooms_ids.discard(yroom.room_id)
+        self.log.info(
+            f"Removed room '{yroom.room_id}' from kernel '{self.kernel_name}'. "
+            f"Total rooms: {len(self._yroom_ids)}"
+        )
+    
+    @property
+    def yroom_manager(self) -> YRoomManager:
+        return self._yroom_manager
+
+    def bind_yroom_manager(self, yroom_manager: YRoomManager):
+        """
+        Binds a reference to the `YRoomManager` singleton to this instance. This
+        method must be called before adding a room.
         """
-        self._yrooms.discard(yroom)
+        self._yroom_manager = yroom_manager
 
 
 AbstractDocumentAwareKernelClient.register(DocumentAwareKernelClient)

jupyter_server_documents/session_manager.py

@@ -122,13 +122,16 @@ async def create_session(
         # Get YRoom for this session and store its ID in `self._room_ids`
         yroom = self._init_session_yroom(session_id, real_path)
 
+        # Bind `YRoomManager` to the kernel client
+        kernel_client = self.get_kernel_client(kernel_id)
+        kernel_client.bind_yroom_manager(yroom_manager=self.yroom_manager)
+
         # 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.get_kernel_client(kernel_id)
         await kernel_client.add_yroom(yroom)
         self.log.info(f"Connected yroom {yroom.room_id} to kernel {kernel_id}. yroom: {yroom}")
         return session_model