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/contents/fps_contents/fileid.py

@@ -7,8 +7,6 @@
 from anyio import Path
 from watchfiles import Change, awatch
 
-from jupyverse_api import Singleton
-
 logger = logging.getLogger("contents")
 
 
@@ -30,7 +28,7 @@ def notify(self, change):
         self._event.set()
 
 
-class FileIdManager(metaclass=Singleton):
+class FileIdManager:
     db_path: str
     initialized: asyncio.Event
     watchers: Dict[str, List[Watcher]]

plugins/contents/fps_contents/routes.py

@@ -25,6 +25,8 @@
 
 
 class _Contents(Contents):
+    _file_id_manager: FileIdManager | None = None
+
     async def create_checkpoint(
         self,
         path,
@@ -245,7 +247,9 @@ async def write_content(self, content: Union[SaveContent, Dict]) -> None:
 
     @property
     def file_id_manager(self):
-        return FileIdManager()
+        if self._file_id_manager is None:
+            self._file_id_manager = FileIdManager()
+        return self._file_id_manager
 
 
 def get_available_path(path: Path, sep: str = "") -> Path:

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_app.py

@@ -1,3 +1,5 @@
+import os
+
 import pytest
 from asphalt.core import Context
 from fastapi import APIRouter
@@ -17,7 +19,8 @@
         "/foo",
     ),
 )
-async def test_mount_path(mount_path, unused_tcp_port):
+async def test_mount_path(mount_path, unused_tcp_port, tmp_path):
+    os.chdir(tmp_path)
     components = configure({"app": {"type": "app"}}, {"app": {"mount_path": mount_path}})
 
     async with Context() as ctx, AsyncClient() as http:

tests/test_auth.py

@@ -1,3 +1,5 @@
+import os
+
 import pytest
 from asphalt.core import Context
 from httpx import AsyncClient
@@ -20,7 +22,8 @@
 
 
 @pytest.mark.asyncio
-async def test_kernel_channels_unauthenticated(unused_tcp_port):
+async def test_kernel_channels_unauthenticated(unused_tcp_port, tmp_path):
+    os.chdir(tmp_path)
     async with Context() as ctx:
         await JupyverseComponent(
             components=COMPONENTS,
@@ -35,7 +38,8 @@ async def test_kernel_channels_unauthenticated(unused_tcp_port):
 
 
 @pytest.mark.asyncio
-async def test_kernel_channels_authenticated(unused_tcp_port):
+async def test_kernel_channels_authenticated(unused_tcp_port, tmp_path):
+    os.chdir(tmp_path)
     async with Context() as ctx, AsyncClient() as http:
         await JupyverseComponent(
             components=COMPONENTS,
@@ -52,7 +56,8 @@ async def test_kernel_channels_authenticated(unused_tcp_port):
 
 @pytest.mark.asyncio
 @pytest.mark.parametrize("auth_mode", ("noauth", "token", "user"))
-async def test_root_auth(auth_mode, unused_tcp_port):
+async def test_root_auth(auth_mode, unused_tcp_port, tmp_path):
+    os.chdir(tmp_path)
     components = configure(COMPONENTS, {"auth": {"mode": auth_mode}})
     async with Context() as ctx, AsyncClient() as http:
         await JupyverseComponent(
@@ -72,7 +77,8 @@ async def test_root_auth(auth_mode, unused_tcp_port):
 
 @pytest.mark.asyncio
 @pytest.mark.parametrize("auth_mode", ("noauth",))
-async def test_no_auth(auth_mode, unused_tcp_port):
+async def test_no_auth(auth_mode, unused_tcp_port, tmp_path):
+    os.chdir(tmp_path)
     components = configure(COMPONENTS, {"auth": {"mode": auth_mode}})
     async with Context() as ctx, AsyncClient() as http:
         await JupyverseComponent(
@@ -86,7 +92,8 @@ async def test_no_auth(auth_mode, unused_tcp_port):
 
 @pytest.mark.asyncio
 @pytest.mark.parametrize("auth_mode", ("token",))
-async def test_token_auth(auth_mode, unused_tcp_port):
+async def test_token_auth(auth_mode, unused_tcp_port, tmp_path):
+    os.chdir(tmp_path)
     components = configure(COMPONENTS, {"auth": {"mode": auth_mode}})
     async with Context() as ctx, AsyncClient() as http:
         await JupyverseComponent(
@@ -113,7 +120,8 @@ async def test_token_auth(auth_mode, unused_tcp_port):
         {"admin": ["read"], "foo": ["bar", "baz"]},
     ),
 )
-async def test_permissions(auth_mode, permissions, unused_tcp_port):
+async def test_permissions(auth_mode, permissions, unused_tcp_port, tmp_path):
+    os.chdir(tmp_path)
     components = configure(COMPONENTS, {"auth": {"mode": auth_mode}})
     async with Context() as ctx, AsyncClient() as http:
         await JupyverseComponent(

tests/test_contents.py

@@ -19,7 +19,6 @@
 @pytest.mark.asyncio
 @pytest.mark.parametrize("auth_mode", ("noauth",))
 async def test_tree(auth_mode, tmp_path, unused_tcp_port):
-    prev_dir = os.getcwd()
     os.chdir(tmp_path)
     dname = Path(".")
     expected = []
@@ -81,4 +80,3 @@ async def test_tree(auth_mode, tmp_path, unused_tcp_port):
         sort_content_by_name(actual)
         sort_content_by_name(expected)
         assert actual == expected
-        os.chdir(prev_dir)

tests/test_execute.py

@@ -2,6 +2,7 @@
 import os
 from functools import partial
 from pathlib import Path
+from shutil import copytree, rmtree
 
 import pytest
 from asphalt.core import Context
@@ -27,6 +28,8 @@
     "yjs": {"type": "yjs"},
 }
 
+HERE = Path(__file__).parent
+
 
 class Websocket:
     def __init__(self, websocket, roomid: str):
@@ -57,7 +60,11 @@ async def recv(self) -> bytes:
 
 @pytest.mark.asyncio
 @pytest.mark.parametrize("auth_mode", ("noauth",))
-async def test_execute(auth_mode, unused_tcp_port):
+async def test_execute(auth_mode, unused_tcp_port, tmp_path):
+    os.chdir(tmp_path)
+    if Path("data").exists():
+        rmtree("data")
+    copytree(HERE / "data", "data")
     url = f"http://127.0.0.1:{unused_tcp_port}"
     components = configure(COMPONENTS, {
         "auth": {"mode": auth_mode},
@@ -71,7 +78,7 @@ async def test_execute(auth_mode, unused_tcp_port):
 
         ws_url = url.replace("http", "ws", 1)
         name = "notebook1.ipynb"
-        path = (Path("tests") / "data" / name).as_posix()
+        path = (Path("data") / name).as_posix()
         # create a session to launch a kernel
         response = await http.post(
             f"{url}/api/sessions",
@@ -92,7 +99,8 @@ async def test_execute(auth_mode, unused_tcp_port):
                 "type": "notebook",
             }
         )
-        file_id = response.json()["fileId"]
+        r = response.json()
+        file_id = r["fileId"]
         document_id = f"json:notebook:{file_id}"
         ynb = ydocs["notebook"]()
         def callback(aevent, events, event):