Backport 'Update jupyter_ydoc and pycrdt_websocket dependencies' jupyterlab#367

Author: brichet

jupyter_collaboration/handlers.py

@@ -8,13 +8,13 @@
 import time
 import uuid
 from logging import Logger
-from typing import Any
+from typing import Any, Literal
 
 from jupyter_server.auth import authorized
 from jupyter_server.base.handlers import APIHandler, JupyterHandler
 from jupyter_server.utils import ensure_async
 from jupyter_ydoc import ydocs as YDOCS
-from pycrdt import YMessageType, write_var_uint
+from pycrdt import write_var_uint
 from pycrdt_websocket.websocket_server import YRoom
 from pycrdt_websocket.ystore import BaseYStore
 from tornado import web
@@ -132,6 +132,10 @@ def exception_logger(exception: Exception, log: Logger) -> bool:
                         exception_handler=exception_logger,
                     )
 
+                    if self._room_id == "JupyterLab:globalAwareness":
+                        # Listen for the changes in GlobalAwareness to update users
+                        self.room.awareness.observe(self._on_global_awareness_event)
+
             try:
                 await self._websocket_server.start_room(self.room)
             except Exception as e:
@@ -280,31 +284,6 @@ async def on_message(self, message):
         """
         message_type = message[0]
 
-        if message_type == YMessageType.AWARENESS:
-            # awareness
-            skip = False
-            changes = self.room.awareness.get_changes(message[1:])
-            added_users = changes["added"]
-            removed_users = changes["removed"]
-            for i, user in enumerate(added_users):
-                u = changes["states"][i]
-                if "user" in u:
-                    name = u["user"]["name"]
-                    self._websocket_server.connected_users[user] = name
-                    self.log.debug("Y user joined: %s", name)
-            for user in removed_users:
-                if user in self._websocket_server.connected_users:
-                    name = self._websocket_server.connected_users[user]
-                    del self._websocket_server.connected_users[user]
-                    self.log.debug("Y user left: %s", name)
-            # filter out message depending on changes
-            if skip:
-                self.log.debug(
-                    "Filtered out Y message of type: %s",
-                    YMessageType(message_type).name,
-                )
-                return skip
-
         if message_type == MessageType.CHAT:
             msg = message[2:].decode("utf-8")
 
@@ -395,6 +374,31 @@ async def _clean_room(self) -> None:
                 self._emit(LogLevel.INFO, "clean", "Loader deleted.")
             del self._room_locks[self._room_id]
 
+    def _on_global_awareness_event(
+        self, topic: Literal["change", "update"], changes: tuple[dict[str, Any], Any]
+    ) -> None:
+        """
+        Update the users when the global awareness changes.
+
+            Parameters:
+                topic (str): `"update"` or `"change"` (`"change"` is triggered only if the states are modified).
+                changes (tuple[dict[str, Any], Any]): The changes and the origin of the changes.
+        """
+        if topic != "change":
+            return
+        added_users = changes[0]["added"]
+        removed_users = changes[0]["removed"]
+        for user in added_users:
+            u = self.room.awareness.states[user]
+            if "user" in u:
+                name = u["user"]["name"]
+                self._websocket_server.connected_users[user] = name
+                self.log.debug("Y user joined: %s", name)
+        for user in removed_users:
+            if user in self._websocket_server.connected_users:
+                name = self._websocket_server.connected_users.pop(user)
+                self.log.debug("Y user left: %s", name)
+
     def check_origin(self, origin):
         """
         Check origin

jupyter_collaboration/rooms.py

@@ -41,7 +41,7 @@ def __init__(
         self._file_format: str = file_format
         self._file_type: str = file_type
         self._file: FileLoader = file
-        self._document = YDOCS.get(self._file_type, YFILE)(self.ydoc)
+        self._document = YDOCS.get(self._file_type, YFILE)(self.ydoc, self.awareness)
 
         self._logger = logger
         self._save_delay = save_delay

pyproject.toml

@@ -28,7 +28,7 @@ classifiers = [
 ]
 dependencies = [
     "jupyter_server>=2.0.0,<3.0.0",
-    "jupyter_ydoc>=2.0.0,<3.0.0",
+    "jupyter_ydoc>=2.1.2,<3.0.0",
     "pycrdt",
     "pycrdt-websocket>=0.14.0,<0.15.0",
     "jupyter_events>=0.10.0",