adding global awareness support

Author: Jialin Zhang

jupyter_rtc_core/app.py

@@ -18,7 +18,7 @@ class RtcExtensionApp(ExtensionApp):
         # # 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),
+        (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"""
@@ -53,8 +53,8 @@ def __init__(
         room_id: str,
         log: Logger,
         loop: asyncio.AbstractEventLoop,
-        fileid_manager: BaseFileIdManager,
-        contents_manager: AsyncContentsManager | ContentsManager,
+        fileid_manager: BaseFileIdManager | None = None,
+        contents_manager: AsyncContentsManager | ContentsManager | None = None,
     ):
         # Bind instance attributes
         self.log = log
@@ -65,23 +65,29 @@ 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()
+
+        self.file_api = None
+        self._jupyter_ydoc = None
+        if fileid_manager and contents_manager: 
+            _, 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 +97,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 +123,8 @@ 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
         return self._jupyter_ydoc
 
 
@@ -127,7 +133,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 +162,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 +456,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 +466,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/rooms/yroom_manager.py

@@ -48,13 +48,18 @@ 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,
-                loop=self.loop,
-                fileid_manager=self.fileid_manager,
-                contents_manager=self.contents_manager,
-            )
+            if room_id == "JupyterLab:globalAwareness":
+                yroom = YRoom(room_id=room_id,
+                    log=self.log,
+                    loop=self.loop)
+            else:
+                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:

jupyter_rtc_core/websockets/__init__.py

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

jupyter_rtc_core/websockets/global_awareness_ws.py

@@ -1,4 +1,49 @@
+from __future__ import annotations
+from tornado.httpclient import HTTPError
 from tornado.websocket import WebSocketHandler
+from typing import TYPE_CHECKING
+from ..rooms import YRoomManager
+import logging
+
+if TYPE_CHECKING:
+    from ..rooms import YRoom
 
 class GlobalAwarenessWebsocket(WebSocketHandler):
-    pass
+    yroom: YRoom
+    client_id: str | None
+    # 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:
+        return self.settings["yroom_manager"]
+
+
+    def prepare(self):
+        # Create the YRoom
+        yroom = self.yroom_manager.get_room("JupyterLab:globalAwareness")
+        if not yroom:
+            raise HTTPError(500, f"Unable to initialize Global Awareness Room.")
+        self.yroom = yroom
+    
+
+    def open(self, *_, **__):
+        # Add self as a client to the YRoom.
+        # If `client_id is None`, then the YRoom is being stopped, and this WS
+        # should be closed immediately w/ close code 1001.
+        self.client_id = self.yroom.clients.add(self)
+        if not self.client_id:
+            self.close(code=1001)
+            return
+
+
+    def on_message(self, message: bytes):
+        # Route all messages to the YRoom for processing
+        self.yroom.add_message(self.client_id, message)
+        
+    def on_close(self):
+        if self.client_id:
+            self.log.info(f"Closed Global Awareness Websocket to client '{self.client_id}'.")
+            self.yroom.clients.remove(self.client_id)
+

jupyter_rtc_core/websockets/yroom_ws.py

@@ -40,11 +40,6 @@ 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)
@@ -53,11 +48,6 @@ def prepare(self):
 
 
     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 +64,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 +73,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)

src/index.ts

@@ -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();