Support YRoom forking

Author: davidbrochart

jupyverse_api/jupyverse_api/yjs/__init__.py

@@ -35,6 +35,13 @@ async def create_roomid(
         ):
             return await self.create_roomid(path, request, response, user)
 
+        @router.put("/api/collaboration/room/{roomid}", status_code=201)
+        async def fork_room(
+            roomid,
+            user: User = Depends(auth.current_user(permissions={"contents": ["read"]})),
+        ):
+            return await self.fork_room(roomid, user)
+
         self.include_router(router)
 
     @abstractmethod
@@ -55,6 +62,14 @@ async def create_roomid(
     ):
         ...
 
+    @abstractmethod
+    async def fork_room(
+        self,
+        roomid,
+        user: User,
+    ):
+        ...
+
     @abstractmethod
     def get_document(
         self,

plugins/yjs/fps_yjs/routes.py

@@ -95,6 +95,26 @@ async def create_roomid(
         res["fileId"] = idx
         return res
 
+    async def fork_room(
+        self,
+        roomid: str,
+        user: User,
+    ):
+        idx = uuid4().hex
+
+        root_room = await self.room_manager.websocket_server.get_room(roomid)
+        update = root_room.ydoc.get_update()
+        new_ydoc = Doc()
+        new_ydoc.apply_update(update)
+        new_room = await self.room_manager.websocket_server.get_room(idx, new_ydoc)
+        root_room.local_clients.add(new_room)
+
+        res = {
+            "sessionId": SERVER_SESSION,
+            "roomId": idx,
+        }
+        return res
+
     def get_document(self, document_id: str) -> YBaseDoc:
         return self.room_manager.documents[document_id]
 
@@ -223,7 +243,7 @@ async def serve(self, websocket: YWebsocket, permissions) -> None:
         await self.websocket_server.started.wait()
         await self.websocket_server.serve(websocket)
 
-        if is_stored_document and not room.clients:
+        if is_stored_document and not room.remote_clients:
             # no client in this room after we disconnect
             self.cleaners[room] = asyncio.create_task(self.maybe_clean_room(room, websocket.path))
 

plugins/yjs/fps_yjs/ywebsocket/websocket_server.py

@@ -151,7 +151,7 @@ async def _serve(self, websocket: Websocket, tg: TaskGroup):
         await self.start_room(room)
         await room.serve(websocket)
 
-        if self.auto_clean_rooms and not room.clients:
+        if self.auto_clean_rooms and not room.remote_clients:
             self.delete_room(room=room)
         tg.cancel_scope.cancel()
 

plugins/yjs/fps_yjs/ywebsocket/yroom.py

@@ -29,7 +29,8 @@
 
 
 class YRoom:
-    clients: list
+    remote_clients: set
+    local_clients: set[YRoom]
     ydoc: Doc
     ystore: BaseYStore | None
     _on_message: Callable[[bytes], Awaitable[bool] | bool] | None
@@ -76,7 +77,8 @@ def __init__(
         self.ready = ready
         self.ystore = ystore
         self.log = log or getLogger(__name__)
-        self.clients = []
+        self.remote_clients = set()
+        self.local_clients = set()
         self._on_message = None
         self._started = None
         self._starting = False
@@ -133,10 +135,15 @@ async def _broadcast_updates(self):
                     return
                 # broadcast internal ydoc's update to all clients, that includes changes from the
                 # clients and changes from the backend (out-of-band changes)
-                for client in self.clients:
-                    self.log.debug("Sending Y update to client with endpoint: %s", client.path)
+                for client in self.local_clients:
+                    client.ydoc.apply_update(update)
+                if self.remote_clients:
                     message = create_update_message(update)
-                    self._task_group.start_soon(client.send, message)
+                    for client in self.remote_clients:
+                        self.log.debug(
+                            "Sending Y update to remote client with endpoint: %s", client.path
+                        )
+                        self._task_group.start_soon(client.send, message)
                 if self.ystore:
                     self.log.debug("Writing Y update to YStore")
                     self._task_group.start_soon(self.ystore.write, update)
@@ -197,7 +204,7 @@ async def serve(self, websocket: Websocket):
             websocket: The WebSocket through which to serve the client.
         """
         async with create_task_group() as tg:
-            self.clients.append(websocket)
+            self.remote_clients.add(websocket)
             await sync(self.ydoc, websocket, self.log)
             try:
                 async for message in websocket:
@@ -224,7 +231,7 @@ async def serve(self, websocket: Websocket):
                             YMessageType.AWARENESS.name,
                             websocket.path,
                         )
-                        for client in self.clients:
+                        for client in self.remote_clients:
                             self.log.debug(
                                 "Sending Y awareness from client with endpoint "
                                 "%s to client with endpoint: %s",
@@ -236,4 +243,4 @@ async def serve(self, websocket: Websocket):
                 self.log.debug("Error serving endpoint: %s", websocket.path, exc_info=e)
 
             # remove this client
-            self.clients = [c for c in self.clients if c != websocket]
+            self.remote_clients.remove(websocket)

tests/test_yjs.py

@@ -0,0 +1,103 @@
+import os
+from asyncio import sleep
+from pathlib import Path
+
+import pytest
+from asphalt.core import Context
+from fps_yjs.ywebsocket import WebsocketProvider
+from httpx import AsyncClient
+from httpx_ws import aconnect_ws
+from pycrdt import Doc, Text
+
+from jupyverse_api.main import JupyverseComponent
+from jupyverse_api.yjs.models import CreateDocumentSession
+
+
+@pytest.mark.asyncio
+async def test_fork_room(tmp_path, unused_tcp_port):
+    prev_dir = os.getcwd()
+    os.chdir(tmp_path)
+
+    path = Path("foo.txt")
+    path.write_text("Hello")
+
+    components = {
+        "app": {"type": "app"},
+        "auth": {"type": "auth", "test": True, "mode": "noauth"},
+        "contents": {"type": "contents"},
+        "frontend": {"type": "frontend"},
+        "yjs": {"type": "yjs"},
+    }
+    async with Context() as ctx, AsyncClient() as http:
+        await JupyverseComponent(
+            components=components,
+            port=unused_tcp_port,
+        ).start(ctx)
+
+        create_document_session = CreateDocumentSession(format="text", type="file")
+        response = await http.put(
+            f"http://127.0.0.1:{unused_tcp_port}/api/collaboration/session/{path}",
+            json=create_document_session.model_dump(),
+        )
+        r = response.json()
+        file_id = r["fileId"]
+
+        # connect to root room
+        async with aconnect_ws(
+            f"http://127.0.0.1:{unused_tcp_port}/api/collaboration/room/text:file:{file_id}"
+        ) as root_ws:
+            # create a root room client
+            root_ydoc = Doc()
+            root_ydoc["source"] = root_ytext = Text()
+            async with WebsocketProvider(root_ydoc, Websocket(root_ws, file_id)):
+                await sleep(0.1)
+                assert str(root_ytext) == "Hello"
+                # fork room
+                response = await http.put(
+                    f"http://127.0.0.1:{unused_tcp_port}/api/collaboration/room/text:file:{file_id}"
+                )
+                r = response.json()
+                fork_room_id = r["roomId"]
+                # connect to forked room
+                async with aconnect_ws(
+                    f"http://127.0.0.1:{unused_tcp_port}/api/collaboration/room/{fork_room_id}"
+                ) as fork_ws:
+                    # create a forked room client
+                    fork_ydoc = Doc()
+                    fork_ydoc["source"] = fork_ytext = Text()
+                    async with WebsocketProvider(fork_ydoc, Websocket(fork_ws, fork_room_id)):
+                        await sleep(0.1)
+                        assert str(fork_ytext) == "Hello"
+                        root_ytext += ", World!"
+                        await sleep(0.1)
+                        assert str(root_ytext) == "Hello, World!"
+                        assert str(fork_ytext) == "Hello, World!"
+
+    os.chdir(prev_dir)
+
+
+class Websocket:
+    def __init__(self, websocket, roomid: str):
+        self.websocket = websocket
+        self.roomid = roomid
+
+    @property
+    def path(self) -> str:
+        return self.roomid
+
+    def __aiter__(self):
+        return self
+
+    async def __anext__(self) -> bytes:
+        try:
+            message = await self.recv()
+        except BaseException:
+            raise StopAsyncIteration()
+        return message
+
+    async def send(self, message: bytes):
+        await self.websocket.send_bytes(message)
+
+    async def recv(self) -> bytes:
+        b = await self.websocket.receive_bytes()
+        return bytes(b)