From 9eed731897790f967a2457be0560db0a3e6f0a77 Mon Sep 17 00:00:00 2001
From: Nicolas Brichet <nicolas.brichet@quantstack.net>
Date: Thu, 10 Oct 2024 10:38:32 +0200
Subject: [PATCH] Update jupyter_ydoc and pycrdt_websocket dependencies (and
 indirectly pycrdt)

---
 .../jupyter_server_ydoc/handlers.py           | 25 -------------------
 .../jupyter_server_ydoc/rooms.py              | 22 +++++++++++++++-
 projects/jupyter-server-ydoc/pyproject.toml   |  4 +--
 3 files changed, 23 insertions(+), 28 deletions(-)

diff --git a/projects/jupyter-server-ydoc/jupyter_server_ydoc/handlers.py b/projects/jupyter-server-ydoc/jupyter_server_ydoc/handlers.py
index f04004fe..8c9f46d4 100644
--- a/projects/jupyter-server-ydoc/jupyter_server_ydoc/handlers.py
+++ b/projects/jupyter-server-ydoc/jupyter_server_ydoc/handlers.py
@@ -286,31 +286,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")
 
diff --git a/projects/jupyter-server-ydoc/jupyter_server_ydoc/rooms.py b/projects/jupyter-server-ydoc/jupyter_server_ydoc/rooms.py
index 0fdb6d8e..692b305c 100644
--- a/projects/jupyter-server-ydoc/jupyter_server_ydoc/rooms.py
+++ b/projects/jupyter-server-ydoc/jupyter_server_ydoc/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._document.path = self._file.path
 
         self._logger = logger
@@ -57,6 +57,9 @@ def __init__(
         self._document.observe(self._on_document_change)
         self._file.observe(self.room_id, self._on_outofband_change, self._on_filepath_change)
 
+        # Listen for awareness changes
+        self.awareness.observe(self._on_awareness_change)
+
     @property
     def file_format(self) -> str:
         """Document file format."""
@@ -316,6 +319,23 @@ async def _maybe_save_document(self, saving_document: asyncio.Task | None) -> No
             self.log.error(msg, exc_info=e)
             self._emit(LogLevel.ERROR, None, msg)
 
+    def _on_awareness_change(self, type: str, changes: tuple[dict[str, Any], Any]):
+        if type != "change":
+            return
+        added_users = changes[0]["added"]
+        removed_users = changes[0]["removed"]
+        for i, user in enumerate(added_users):
+            u = self.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[user]
+                del self._websocket_server.connected_users[user]
+                self.log.debug("Y user left: %s", name)
+
 
 class TransientRoom(YRoom):
     """A Y room for sharing state (e.g. awareness)."""
diff --git a/projects/jupyter-server-ydoc/pyproject.toml b/projects/jupyter-server-ydoc/pyproject.toml
index 3348d07d..7007980e 100644
--- a/projects/jupyter-server-ydoc/pyproject.toml
+++ b/projects/jupyter-server-ydoc/pyproject.toml
@@ -29,9 +29,9 @@ authors = [
 ]
 dependencies = [
     "jupyter_server>=2.11.1,<3.0.0",
-    "jupyter_ydoc>=2.0.0,<4.0.0",
+    "jupyter_ydoc>=2.1.2,<4.0.0",
     "pycrdt",
-    "pycrdt-websocket>=0.14.2,<0.15.0",
+    "pycrdt-websocket>=0.15.0,<0.16.0",
     "jupyter_events>=0.10.0",
     "jupyter_server_fileid>=0.7.0,<1",
     "jsonschema>=4.18.0"