From 35ee5b5354cd5d2cc4ee3f4e003260d52a4aaa84 Mon Sep 17 00:00:00 2001
From: David Brochart <david.brochart@gmail.com>
Date: Wed, 28 Jun 2023 18:21:11 +0200
Subject: [PATCH] Support ypy-websocket v0.12 (#320)

---
 plugins/yjs/fps_yjs/routes.py | 11 ++++++++---
 plugins/yjs/pyproject.toml    |  2 +-
 tests/test_server.py          |  7 ++++---
 3 files changed, 13 insertions(+), 7 deletions(-)

diff --git a/plugins/yjs/fps_yjs/routes.py b/plugins/yjs/fps_yjs/routes.py
index a016952a..148c5045 100644
--- a/plugins/yjs/fps_yjs/routes.py
+++ b/plugins/yjs/fps_yjs/routes.py
@@ -148,6 +148,7 @@ def __init__(self, contents: Contents):
         self.cleaners = {}  # a dictionary of room:task
         self.last_modified = {}  # a dictionary of file_id:last_modification_date
         self.websocket_server = JupyterWebsocketServer(rooms_ready=False, auto_clean_rooms=False)
+        self.websocket_server_task = asyncio.create_task(self.websocket_server.start())
         self.lock = asyncio.Lock()
 
     def stop(self):
@@ -157,9 +158,10 @@ def stop(self):
             saver.cancel()
         for cleaner in self.cleaners.values():
             cleaner.cancel()
+        self.websocket_server.stop()
 
     async def serve(self, websocket: YpyWebsocket, permissions) -> None:
-        room = self.websocket_server.get_room(websocket.path)
+        room = await self.websocket_server.get_room(websocket.path)
         can_write = permissions is None or "write" in permissions.get("yjs", [])
         room.on_message = partial(self.filter_message, can_write)
         is_stored_document = websocket.path.count(":") >= 2
@@ -209,6 +211,7 @@ async def serve(self, websocket: YpyWebsocket, permissions) -> None:
                             self.watch_file(file_format, file_id, document)
                         )
 
+        await self.websocket_server.started.wait()
         await self.websocket_server.serve(websocket)
 
         if is_stored_document and not room.clients:
@@ -354,7 +357,7 @@ async def maybe_clean_room(self, room, ws_path: str) -> None:
 
 
 class JupyterWebsocketServer(WebsocketServer):
-    def get_room(self, ws_path: str) -> YRoom:
+    async def get_room(self, ws_path: str) -> YRoom:
         if ws_path not in self.rooms:
             if ws_path.count(":") >= 2:
                 # it is a stored document (e.g. a notebook)
@@ -365,4 +368,6 @@ def get_room(self, ws_path: str) -> YRoom:
             else:
                 # it is a transient document (e.g. awareness)
                 self.rooms[ws_path] = YRoom()
-        return self.rooms[ws_path]
+        room = self.rooms[ws_path]
+        await self.start_room(room)
+        return room
diff --git a/plugins/yjs/pyproject.toml b/plugins/yjs/pyproject.toml
index 715cd1fd..e6bd2d81 100644
--- a/plugins/yjs/pyproject.toml
+++ b/plugins/yjs/pyproject.toml
@@ -9,7 +9,7 @@ keywords = [ "jupyter", "server", "fastapi", "plugins" ]
 requires-python = ">=3.8"
 dependencies = [
     "jupyter_ydoc >=1,<2",
-    "ypy-websocket >=0.8.2,<0.11.0",
+    "ypy-websocket >=0.12.1,<0.13.0",
     "y-py >=0.6.0,<0.7.0",
     "jupyverse-api >=0.1.2,<1",
 ]
diff --git a/tests/test_server.py b/tests/test_server.py
index 8286dea0..07993dbc 100644
--- a/tests/test_server.py
+++ b/tests/test_server.py
@@ -79,10 +79,11 @@ async def test_rest_api(start_jupyverse):
     )
     file_id = response.json()["fileId"]
     document_id = f"json:notebook:{file_id}"
-    async with connect(f"{ws_url}/api/collaboration/room/{document_id}") as websocket:
+    ydoc = Y.YDoc()
+    async with connect(
+        f"{ws_url}/api/collaboration/room/{document_id}"
+    ) as websocket, WebsocketProvider(ydoc, websocket):
         # connect to the shared notebook document
-        ydoc = Y.YDoc()
-        WebsocketProvider(ydoc, websocket)
         # wait for file to be loaded and Y model to be created in server and client
         await asyncio.sleep(0.5)
         # execute notebook