revert changes to kernel client

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,9 +16,6 @@
 
 from .kernel_client_abc import AbstractDocumentAwareKernelClient
 
-if t.TYPE_CHECKING:
-    from jupyter_server_documents.rooms.yroom_manager import YRoomManager
-
 
 class DocumentAwareKernelClient(AsyncKernelClient):
     """
@@ -42,6 +39,11 @@ 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
@@ -56,24 +58,6 @@ 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.
@@ -302,63 +286,18 @@ 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, 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.
+        Register a YRoom with this kernel client. YRooms will
+        intercept display and kernel status messages.
         """
-        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)}"
-        )
+        self._yrooms.add(yroom)
 
     async def remove_yroom(self, yroom: YRoom):
         """
-        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.
+        De-register a YRoom from handling kernel client messages.
         """
-        self._yroom_manager = yroom_manager
+        self._yrooms.discard(yroom)
 
 
 AbstractDocumentAwareKernelClient.register(DocumentAwareKernelClient)

jupyter_server_documents/session_manager.py

@@ -122,16 +122,13 @@ 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