Add support for global awareness (#63)

Co-authored-by: Jialin Zhang <[email protected]>

Author: jzhang20133

jupyter_rtc_core/app.py

@@ -4,7 +4,7 @@
 
 from traitlets import Instance, Type
 from .handlers import RouteHandler, FileIDIndexHandler
-from .websockets import GlobalAwarenessWebsocket, YRoomWebsocket
+from .websockets import YRoomWebsocket
 from .rooms.yroom_manager import YRoomManager
 from .outputs import OutputsManager, outputs_handlers
 
@@ -17,8 +17,6 @@ class RtcExtensionApp(ExtensionApp):
         # dummy handler that verifies the server extension is installed;
         # # this can be deleted prior to initial release.
         (r"jupyter-rtc-core/get-example/?", RouteHandler),
-        # global awareness websocket
-        # (r"api/collaboration/room/JupyterLab:globalAwareness/?", GlobalAwarenessWebsocket),
         # # ydoc websocket
         (r"api/collaboration/room/(.*)", YRoomWebsocket),
         # # handler that just adds compatibility with Jupyter Collaboration's frontend

jupyter_rtc_core/rooms/yroom.py

@@ -29,7 +29,7 @@ class YRoom:
     room_id := "{file_type}:{file_format}:{file_id}"
     """
 
-    _jupyter_ydoc: YBaseDoc
+    _jupyter_ydoc: YBaseDoc | None
     """JupyterYDoc"""
     _ydoc: pycrdt.Doc
     """Ydoc"""
@@ -65,23 +65,34 @@ def __init__(
         self._client_group = YjsClientGroup(room_id=room_id, log=self.log, loop=self._loop)
         self._ydoc = pycrdt.Doc()
         self._awareness = pycrdt.Awareness(ydoc=self._ydoc)
-        _, file_type, _ = self.room_id.split(":")
-        JupyterYDocClass = cast(
-            type[YBaseDoc],
-            jupyter_ydoc_classes.get(file_type, jupyter_ydoc_classes["file"])
-        )
-        self._jupyter_ydoc = JupyterYDocClass(ydoc=self._ydoc, awareness=self._awareness)
-
-        # Initialize YRoomFileAPI and begin loading content
-        self.file_api = YRoomFileAPI(
-            room_id=self.room_id,
-            jupyter_ydoc=self._jupyter_ydoc,
-            log=self.log,
-            loop=self._loop,
-            fileid_manager=fileid_manager,
-            contents_manager=contents_manager
-        )
-        self.file_api.load_ydoc_content()
+
+        # If this room is providing global awareness, set
+        # `file_api` and `jupyter_ydoc` to `None` as this room
+        # will never read/write via the `ContentsManager`.
+        if self.room_id == "JupyterLab:globalAwareness":
+            self.file_api = None
+            self._jupyter_ydoc = None
+        else:
+            # Otherwise, initialize `jupyter_ydoc` and `file_api`
+            _, file_type, _ = self.room_id.split(":")
+            JupyterYDocClass = cast(
+                type[YBaseDoc],
+                jupyter_ydoc_classes.get(file_type, jupyter_ydoc_classes["file"])
+            )
+            self._jupyter_ydoc = JupyterYDocClass(ydoc=self._ydoc, awareness=self._awareness)
+
+            # Initialize YRoomFileAPI and begin loading content
+            self.file_api = YRoomFileAPI(
+                room_id=self.room_id,
+                jupyter_ydoc=self._jupyter_ydoc,
+                log=self.log,
+                loop=self._loop,
+                fileid_manager=fileid_manager,
+                contents_manager=contents_manager
+            )
+            self.file_api.load_ydoc_content()
+            self._jupyter_ydoc.observe(self._on_jupyter_ydoc_update)
+        
 
         # Start observers on `self.ydoc` and `self.awareness` to ensure new
         # updates are broadcast to all clients and saved to disk.
@@ -91,7 +102,6 @@ def __init__(
         self._ydoc_subscription = self._ydoc.observe(
             self._on_ydoc_update
         )
-        self._jupyter_ydoc.observe(self._on_jupyter_ydoc_update)
 
         # Initialize message queue and start background task that routes new
         # messages in the message queue to the appropriate handler method.
@@ -118,7 +128,12 @@ async def get_jupyter_ydoc(self):
         (`jupyter_ydoc.ybasedoc.YBaseDoc`) after waiting for its content to be
         loaded from the ContentsManager.
         """
-        await self.file_api.ydoc_content_loaded
+        if self.file_api:
+            await self.file_api.ydoc_content_loaded
+        if self.room_id == "JupyterLab:globalAwareness":
+            message = "There is no Jupyter ydoc for global awareness scenario"
+            self.log.error(message)
+            raise Exception(message)
         return self._jupyter_ydoc
 
 
@@ -127,7 +142,8 @@ async def get_ydoc(self):
         Returns a reference to the room's YDoc (`pycrdt.Doc`) after
         waiting for its content to be loaded from the ContentsManager.
         """
-        await self.file_api.ydoc_content_loaded
+        if self.file_api:
+            await self.file_api.ydoc_content_loaded
         return self._ydoc
 
 
@@ -155,7 +171,8 @@ async def _process_message_queue(self) -> None:
         """
         # Wait for content to be loaded before processing any messages in the
         # message queue
-        await self.file_api.ydoc_content_loaded
+        if self.file_api:
+            await self.file_api.ydoc_content_loaded
 
         # Begin processing messages from the message queue
         while True:
@@ -448,7 +465,8 @@ async def stop(self) -> None:
         # Remove all observers, as updates no longer need to be broadcast
         self._ydoc.unobserve(self._ydoc_subscription)
         self._awareness.unobserve(self._awareness_subscription)
-        self._jupyter_ydoc.unobserve()
+        if self._jupyter_ydoc:
+            self._jupyter_ydoc.unobserve()
 
         # Finish processing all messages, then stop the queue to end the
         # `_process_message_queue()` background task.
@@ -457,4 +475,5 @@ async def stop(self) -> None:
 
         # Finally, stop FileAPI and return. This saves the final content of the
         # JupyterYDoc in the process.
-        await self.file_api.stop()
+        if self.file_api:
+            await self.file_api.stop()

jupyter_rtc_core/websockets/__init__.py

@@ -1,3 +1,2 @@
-from .global_awareness_ws import GlobalAwarenessWebsocket
 from .clients import YjsClient, YjsClientGroup
 from .yroom_ws import YRoomWebsocket

jupyter_rtc_core/websockets/global_awareness_ws.py

@@ -40,24 +40,15 @@ def prepare(self):
         request_path: str = self.request.path
         self.room_id = request_path.strip("/").split("/")[-1]
 
-        # TODO: remove this once globalawareness is implemented
-        if self.room_id == "JupyterLab:globalAwareness":
-            self.close(1011)
-            return
-
         # 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}'.")
+        if self.room_id != "JupyterLab:globalAwareness":
+            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, *_, **__):
-        # TODO: remove this later
-        if self.room_id == "JupyterLab:globalAwareness":
-            self.close(1011)
-            return
-
         # Create the YRoom
         yroom = self.yroom_manager.get_room(self.room_id)
         if not yroom:
@@ -74,10 +65,6 @@ def open(self, *_, **__):
 
 
     def on_message(self, message: bytes):
-        # TODO: remove this later
-        if self.room_id == "JupyterLab:globalAwareness":
-            return
-
         if not self.client_id:
             self.close(code=1001)
             return
@@ -87,10 +74,6 @@ def on_message(self, message: bytes):
 
 
     def on_close(self):
-        # TODO: remove this later
-        if self.room_id == "JupyterLab:globalAwareness":
-            return
-
         if self.client_id:
             self.log.info(f"Closed Websocket to client '{self.client_id}'.")
             self.yroom.clients.remove(self.client_id)

jupyter_rtc_core/websockets/yroom_ws.py

@@ -46,6 +46,9 @@ import { IAwareness } from '@jupyter/ydoc';
 import {
   AwarenessKernelStatus
 } from './kernelstatus';
+import { ServerConnection } from '@jupyterlab/services';
+import { WebSocketAwarenessProvider } from './docprovider/awareness';
+import { URLExt } from '@jupyterlab/coreutils';
 /**
  * Initialization data for the @jupyter/rtc-core extension.
  */
@@ -102,15 +105,15 @@ export const rtcGlobalAwarenessPlugin: JupyterFrontEndPlugin<IAwareness> = {
     const awareness = new Awareness(ydoc);
 
     // TODO: Uncomment once global awareness is working
-    /*const server = ServerConnection.makeSettings();
+    const server = ServerConnection.makeSettings();
     const url = URLExt.join(server.wsUrl, 'api/collaboration/room');
 
     new WebSocketAwarenessProvider({
       url: url,
       roomID: 'JupyterLab:globalAwareness',
       awareness: awareness,
       user: user
-    });*/
+    });
 
     state.changed.connect(async () => {
       const data: any = await state.toJSON();