Add YRoomWebsocket and YRoomManager (#27)

* add yroom websocket

* fix add_message() call and impl prepare()

* add YRoomManager

* fix runtime errors caused by changes

* fix import in kernels module to allow JL to start

* load fileid_manager lazily to allow JL to start

Author: dlqqq

jupyter_rtc_core/app.py

@@ -1,8 +1,10 @@
 from jupyter_server.extension.application import ExtensionApp
 from traitlets.config import Config
+import asyncio
 
 from .handlers import RouteHandler
 from .websockets import GlobalAwarenessWebsocket, YRoomWebsocket
+from .rooms import YRoomManager
 
 class RtcExtensionApp(ExtensionApp):
     name = "jupyter_rtc_core"
@@ -21,6 +23,25 @@ class RtcExtensionApp(ExtensionApp):
 
     def initialize(self):
         super().initialize()
+
+
+    def initialize_settings(self):
+        # Get YRoomManager arguments from server extension context.
+        # We cannot access the 'file_id_manager' key immediately because server
+        # extensions initialize in alphabetical order. 'jupyter_rtc_core' <
+        # 'jupyter_server_fileid'.
+        get_fileid_manager = lambda: self.settings["file_id_manager"]
+        contents_manager = self.serverapp.contents_manager
+        loop = asyncio.get_event_loop_policy().get_event_loop()
+        log = self.log
+
+        # Initialize YRoomManager
+        self.settings["yroom_manager"] = YRoomManager(
+            get_fileid_manager=get_fileid_manager,
+            contents_manager=contents_manager,
+            loop=loop,
+            log=log
+        )
 
 
     def _link_jupyter_server_extension(self, server_app):

jupyter_rtc_core/kernels/websocket_connection.py

@@ -14,7 +14,7 @@
 from jupyter_server.services.kernels.connection.base import (
     BaseKernelWebsocketConnection,
 )
-from .states import LIFECYCLE_DEAD_STATES
+from .states import LifecycleStates
 from jupyter_server.services.kernels.connection.base import deserialize_msg_from_ws_v1, serialize_msg_to_ws_v1
 from jupyter_client.session import Session
 

jupyter_rtc_core/rooms/yroom.py

@@ -23,6 +23,7 @@ class YRoom:
     """Log object"""
     room_id: str
     """Room Id"""
+
     _jupyter_ydoc: YBaseDoc
     """JupyterYDoc"""
     _ydoc: pycrdt.Doc
@@ -35,6 +36,10 @@ class YRoom:
     """Client group to manage synced and desynced clients"""
     _message_queue: asyncio.Queue[Tuple[str, bytes]]
     """A message queue per room to keep websocket messages in order"""
+    _awareness_subscription: pycrdt.Subscription
+    """Subscription to awareness changes."""
+    _ydoc_subscription: pycrdt.Subscription
+    """Subscription to YDoc changes."""
 
 
     def __init__(
@@ -74,8 +79,12 @@ def __init__(
 
         # Start observers on `self.ydoc` and `self.awareness` to ensure new
         # updates are broadcast to all clients and saved to disk.
-        self._awareness.observe(self.send_server_awareness)
-        self._ydoc.observe(lambda event: self.write_sync_update(event.update))
+        self._awareness_subscription = self._awareness.observe(
+            self.send_server_awareness
+        )
+        self._ydoc_subscription = self._ydoc.observe(
+            lambda event: self.write_sync_update(event.update)
+        )
 
         # Initialize message queue and start background task that routes new
         # messages in the message queue to the appropriate handler method.
@@ -368,6 +377,13 @@ def send_server_awareness(self, type: str, changes: tuple[dict[str, Any], Any])
         self._broadcast_message(message, "AwarenessUpdate")
 
     def stop(self) -> None:
-        # TODO: requires YRoomLoader implementation
+        """
+        Stop the YRoom.
+
+        TODO: stop file API & stop the message processing loop
+        """
+        self._ydoc.unobserve(self._ydoc_subscription)
+        self._awareness.unobserve(self._awareness_subscription)
+
         return
 

jupyter_rtc_core/rooms/yroom_manager.py

@@ -1,21 +1,79 @@
-"""
-WIP.
-
-This file just contains interfaces to be filled out later.
-"""
-
+from __future__ import annotations
 from .yroom import YRoom
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    import asyncio
+    import logging
+    from typing import Callable
+    from jupyter_server_fileid.manager import BaseFileIdManager
+    from jupyter_server.services.contents.manager import AsyncContentsManager, ContentsManager
 
 class YRoomManager:
     _rooms_by_id: dict[str, YRoom]
 
-    def __init__(self):
+    def __init__(
+        self,
+        *,
+        get_fileid_manager: Callable[[], BaseFileIdManager],
+        contents_manager: AsyncContentsManager | ContentsManager,
+        loop: asyncio.AbstractEventLoop,
+        log: logging.Logger,
+    ):
+        # Bind instance attributes
+        self._get_fileid_manager = get_fileid_manager
+        self.contents_manager = contents_manager
+        self.loop = loop
+        self.log = log
+
+        # Initialize dictionary of YRooms, keyed by room ID
         self._rooms_by_id = {}
 
+
+    @property
+    def fileid_manager(self) -> BaseFileIdManager:
+        return self._get_fileid_manager()
+    
+
     def get_room(self, room_id: str) -> YRoom | None:
-        # TODO
-        return None
+        """
+        Retrieves a YRoom given a room ID. If the YRoom does not exist, this
+        method will initialize a new YRoom.
+        """
+
+        # If room exists, then return it immediately
+        if room_id in self._rooms_by_id:
+            return self._rooms_by_id[room_id]
+        
+        # Otherwise, create a new room
+        try:
+            yroom = YRoom(
+                room_id=room_id,
+                log=self.log,
+                loop=self.loop,
+                fileid_manager=self.fileid_manager,
+                contents_manager=self.contents_manager,
+            )
+            self._rooms_by_id[room_id] = yroom
+            return yroom
+        except Exception as e:
+            self.log.error(
+                f"Unable to initialize YRoom '{room_id}'.",
+                exc_info=True
+            )
+            return None
+        
+        
+    def delete_room(self, room_id: str) -> None:
+        """
+        Deletes a YRoom given a room ID.
+        
+        TODO: finish implementing YRoom.stop(), and delete empty rooms w/ no
+        live kernels automatically in a background task.
+        """
+        yroom = self._rooms_by_id.get(room_id, None)
+        if not yroom:
+            return
 
-    def delete_room(self, room: YRoom) -> None:
-        # TODO
-        return
+        yroom.stop()
+        del self._rooms_by_id[room_id]

jupyter_rtc_core/websockets/clients.py

@@ -9,7 +9,6 @@
 import datetime
 from logging import Logger
 from typing import TYPE_CHECKING
-from dataclasses import dataclass
 import uuid
 import asyncio
 

jupyter_rtc_core/websockets/yroom_ws.py

@@ -1,17 +1,54 @@
-"""
-WIP.
-
-This file just contains an interface to be filled out later.
-"""
-
+from __future__ import annotations
+from tornado.httpclient import HTTPError
 from tornado.websocket import WebSocketHandler
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from jupyter_server_fileid.manager import BaseFileIdManager
+    from ..rooms import YRoom, YRoomManager
 
 class YRoomWebsocket(WebSocketHandler):
+    yroom: YRoom
+    room_id: str
+    client_id: str
+
+    @property
+    def yroom_manager(self) -> YRoomManager:
+        return self.settings["yroom_manager"]
+    
+
+    @property
+    def fileid_manager(self) -> BaseFileIdManager:
+        return self.settings["file_id_manager"]
+
+
+    def prepare(self):
+        # Bind `room_id` attribute
+        request_path: str = self.request.path
+        self.room_id = request_path.strip("/").split("/")[-1]
+
+        # Verify the file ID contained in the room ID points to a valid file.
+        fileid = self.room_id.split(":")[-1]
+        path = self.fileid_manager.get_path(fileid)
+        if not path:
+            raise HTTPError(404, f"No file with ID '{fileid}'.")
+    
+
     def open(self):
-        print("WebSocket opened")
+        # Create the YRoom
+        yroom = self.yroom_manager.get_room(self.room_id)
+        if not yroom:
+            raise HTTPError(500, f"Unable to initialize YRoom '{self.room_id}'.")
+        self.yroom = yroom
+
+        # Add self as a client to the YRoom
+        self.client_id = self.yroom.clients.add(self)
+
+
+    def on_message(self, message: bytes):
+        # Route all messages to the YRoom for processing
+        self.yroom.add_message(self.client_id, message)
 
-    def on_message(self, message):
-        self.write_message(u"You said: " + message)
 
     def on_close(self):
-        print("WebSocket closed")
+        self.yroom.clients.remove(self.client_id)