fix Y protocol implementation, set binary=True when writing to WS handler

Author: dlqqq

jupyter_rtc_core/app.py

@@ -45,24 +45,19 @@ def initialize_settings(self):
         # We cannot access the 'file_id_manager' key immediately because server
         # extensions initialize in alphabetical order. 'jupyter_rtc_core' <
         # 'jupyter_server_fileid'.
-        # def get_fileid_manager():
-        #     self.log.info("IN GETTER")
-        #     for k, v in self.settings.items():
-        #         print(f"{k}: {v}")
-        #     print(len(self.settings.items()))
-        #     print(id(self.settings))
-        #     return self.settings["file_id_manager"]
-        # contents_manager = self.serverapp.contents_manager
-        # loop = asyncio.get_event_loop_policy().get_event_loop()
-        # log = self.log
+        def get_fileid_manager():
+            return self.serverapp.web_app.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
-        # )
+        # Initialize YRoomManager
+        self.settings["yroom_manager"] = YRoomManager(
+            get_fileid_manager=get_fileid_manager,
+            contents_manager=contents_manager,
+            loop=loop,
+            log=log
+        )
         pass
 
 

jupyter_rtc_core/rooms/yroom.py

@@ -25,7 +25,7 @@ class YRoom:
     """
     The ID of the room. This is a composite ID following the format:
 
-    room_id := "{file_type}:{file_format}
+    room_id := "{file_type}:{file_format}:{file_id}"
     """
 
     _jupyter_ydoc: YBaseDoc
@@ -95,6 +95,9 @@ def __init__(
         # messages in the message queue to the appropriate handler method.
         self._message_queue = asyncio.Queue()
         self._loop.create_task(self._on_new_message())
+
+        # Log notification that room is ready
+        self.log.info(f"Room '{self.room_id}' initialized.")
 
 
     @property
@@ -156,13 +159,15 @@ async def _on_new_message(self) -> None:
         while True:
             try: 
                 client_id, message = await self._message_queue.get()
+                self.log.info(f"HANDLING NEW MESSAGE FROM '{client_id}'")
+                self.log.info(f"Message: {message}")
             except asyncio.QueueShutDown:
                 break
 
             # Handle Awareness messages
             message_type = message[0]
             if message_type == YMessageType.AWARENESS:
-                self.handle_awareness_update(client_id, message[1:])
+                self.handle_awareness_update(client_id, message)
                 continue
 
             # Handle Sync messages
@@ -196,6 +201,8 @@ def handle_sync_step1(self, client_id: str, message: bytes) -> None:
         - Sending the reply to the client over WS, and
         - Sending a new SyncStep1 message immediately after.
         """
+        self.log.info(f"Handling SS1 message from client '{client_id}'.")
+
         # Mark client as desynced
         new_client = self.clients.get(client_id)
         self.clients.mark_desynced(client_id)
@@ -218,7 +225,7 @@ def handle_sync_step1(self, client_id: str, message: bytes) -> None:
         try:
             # TODO: remove the assert once websocket is made required
             assert isinstance(new_client.websocket, WebSocketHandler)
-            new_client.websocket.write_message(sync_step2_message)
+            new_client.websocket.write_message(sync_step2_message, binary=True)
         except Exception as e:
             self.log.error(
                 "An exception occurred when writing the SyncStep2 reply "
@@ -228,18 +235,21 @@ def handle_sync_step1(self, client_id: str, message: bytes) -> None:
             return
 
         self.clients.mark_synced(client_id)
+        self.log.info(f"Sent SS2 reply to client '{client_id}'.")
 
         # Send SyncStep1 message
         try:
             assert isinstance(new_client.websocket, WebSocketHandler)
             sync_step1_message = pycrdt.create_sync_message(self._ydoc)
-            new_client.websocket.write_message(sync_step1_message)
+            new_client.websocket.write_message(sync_step1_message, binary=True)
         except Exception as e:
             self.log.error(
                 "An exception occurred when writing a SyncStep1 message "
                 f"to newly-synced client '{new_client.id}':"
             )
             self.log.exception(e)
+        self.log.info(f"Sent SS1 message to client '{client_id}'.")
+        self.log.info(f"Message: {sync_step1_message}")
 
 
     def handle_sync_step2(self, client_id: str, message: bytes) -> None:
@@ -251,6 +261,7 @@ def handle_sync_step2(self, client_id: str, message: bytes) -> None:
         clients after this method is called via the `self.write_sync_update()`
         observer.
         """
+        self.log.info("HANDLING SS2 MESSAGE")
         try:
             message_payload = message[1:]
             pycrdt.handle_sync_message(message_payload, self._ydoc)
@@ -271,6 +282,7 @@ def handle_sync_update(self, client_id: str, message: bytes) -> None:
         clients after this method is called via the `self.write_sync_update()`
         observer.
         """
+        self.log.info("HANDLING SYNCUPDATE")
         # Remove client and kill websocket if received SyncUpdate when client is desynced
         if self._should_ignore_update(client_id, "SyncUpdate"):
             self.log.error(f"Should not receive SyncUpdate message when double handshake is not completed for client '{client_id}' and room '{self.room_id}'")
@@ -310,6 +322,7 @@ def write_sync_update(self, message_payload: bytes, client_id: str | None = None
 
 
     def handle_awareness_update(self, client_id: str, message: bytes) -> None:
+        self.log.info("HANDLING AWARENESS UPDATE")
         # Apply the AwarenessUpdate message
         try:
             message_payload = pycrdt.read_message(message[1:])
@@ -356,7 +369,7 @@ def _broadcast_message(self, message: bytes, message_type: Literal['AwarenessUpd
             try:
                 # TODO: remove this assertion once websocket is made required
                 assert isinstance(client.websocket, WebSocketHandler)
-                client.websocket.write_message(message)
+                client.websocket.write_message(message, binary=True)
             except Exception as e:
                 self.log.warning(
                     f"An exception occurred when broadcasting a "

jupyter_rtc_core/rooms/yroom_file_api.py

@@ -33,6 +33,7 @@ class YRoomFileAPI:
 
     # See `filemanager.py` in `jupyter_server` for references on supported file
     # formats & file types.
+    room_id: str
     file_format: Literal["text", "base64"]
     file_type: Literal["file", "notebook"]
     file_id: str
@@ -57,6 +58,7 @@ def __init__(
         loop: asyncio.AbstractEventLoop
     ):
         # Bind instance attributes
+        self.room_id = room_id
         self.file_format, self.file_type, self.file_id = room_id.split(":")
         self.jupyter_ydoc = jupyter_ydoc
         self.log = log
@@ -113,6 +115,8 @@ def load_ydoc_content(self) -> None:
         # Otherwise, set loading to `True` and start the loading task.
         if self._ydoc_content_loaded.is_set() or self._ydoc_content_loading:
             return
+        
+        self.log.info(f"Loading content for room ID '{self.room_id}'.")
         self._ydoc_content_loading = True
         self._loop.create_task(self._load_ydoc_content())
 
@@ -134,6 +138,8 @@ async def _load_ydoc_content(self) -> None:
         # Also set loading to `False` for consistency
         self._ydoc_content_loaded.set()
         self._ydoc_content_loading = False
+        self.log.info(f"Loaded content for room ID '{self.room_id}'.")
+        self.log.info(f"Content: {self.jupyter_ydoc.source}")
 
 
     def schedule_save(self) -> None:

jupyter_rtc_core/rooms/yroom_manager.py

@@ -8,7 +8,6 @@
 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
 
@@ -18,20 +17,25 @@ class YRoomManager():
     def __init__(
         self,
         *,
-        fileid_manager: BaseFileIdManager,
+        get_fileid_manager: callable[[], BaseFileIdManager],
         contents_manager: AsyncContentsManager | ContentsManager,
         loop: asyncio.AbstractEventLoop,
         log: logging.Logger,
     ):
         # Bind instance attributes
-        self.fileid_manager = fileid_manager
+        self._get_fileid_manager = get_fileid_manager
         self.contents_manager = contents_manager
         self.loop = loop
         self.log = log
         self._rooms_by_id = {}
         # Initialize dictionary of YRooms, keyed by room ID
 
 
+    @property
+    def fileid_manager(self) -> BaseFileIdManager:
+        return self._get_fileid_manager()
+
+
     def get_room(self, room_id: str) -> YRoom | None:
         """
         Retrieves a YRoom given a room ID. If the YRoom does not exist, this
@@ -44,6 +48,7 @@ def get_room(self, room_id: str) -> YRoom | None:
 
         # Otherwise, create a new room
         try:
+            self.log.info(f"Initializing room '{room_id}'.")
             yroom = YRoom(
                 room_id=room_id,
                 log=self.log,

jupyter_rtc_core/websockets/clients.py

@@ -77,7 +77,7 @@ def __init__(self, *, room_id: str, log: Logger, loop: asyncio.AbstractEventLoop
         self.desynced: dict[str, YjsClient] = {}
         self.log = log
         self.loop = loop
-        self.loop.create_task(self._clean_desynced())
+        # self.loop.create_task(self._clean_desynced())
         self._poll_interval_seconds = poll_interval_seconds
         self.desynced_timeout_seconds = desynced_timeout_seconds
 

jupyter_rtc_core/websockets/yroom_ws.py

@@ -15,21 +15,16 @@ class YRoomWebsocket(WebSocketHandler):
     yroom: YRoom
     room_id: str
     client_id: str
+    # TODO: change this. we should pass `self.log` from our
+    # `ExtensionApp` to log messages w/ "RtcCoreExtension" prefix
+    log = logging.Logger("TEMP")
+
 
     @property
     def yroom_manager(self) -> YRoomManager:
-        if "yroom_manager" not in self.settings:
-            self.settings["yroom_manager"] = YRoomManager(
-                fileid_manager=self.fileid_manager,
-                contents_manager=self.contents_manager,
-                loop=asyncio.get_event_loop_policy().get_event_loop(),
-                # TODO: change this. we should pass `self.log` from our
-                # `ExtensionApp` to log messages w/ "RtcCoreExtension" prefix
-                log=logging.Logger("TEMP")
-
-            )
         return self.settings["yroom_manager"]
 
+
     @property
     def fileid_manager(self) -> BaseFileIdManager:
         return self.settings["file_id_manager"]
@@ -69,9 +64,9 @@ def open(self, *_, **__):
 
     def on_message(self, message: bytes):
         # Route all messages to the YRoom for processing
-        print(message)
         self.yroom.add_message(self.client_id, message)
 
 
     def on_close(self):
+        self.log.info("WEBSOCKET CLOSED")
         self.yroom.clients.remove(self.client_id)