From d8a76ef3f7a4974f45c5cf57f0a73fb51d74c4dc Mon Sep 17 00:00:00 2001
From: David Brochart <david.brochart@gmail.com>
Date: Wed, 2 Oct 2024 16:30:03 +0200
Subject: [PATCH] Add Yjs plugin configuration

---
 jupyverse_api/jupyverse_api/yjs/__init__.py |  7 ++++++-
 plugins/yjs/fps_yjs/main.py                 |  9 +++++++--
 plugins/yjs/fps_yjs/routes.py               | 14 ++++++++------
 3 files changed, 21 insertions(+), 9 deletions(-)

diff --git a/jupyverse_api/jupyverse_api/yjs/__init__.py b/jupyverse_api/jupyverse_api/yjs/__init__.py
index 7a8dfa5b..871a01ce 100644
--- a/jupyverse_api/jupyverse_api/yjs/__init__.py
+++ b/jupyverse_api/jupyverse_api/yjs/__init__.py
@@ -3,7 +3,7 @@
 
 from fastapi import APIRouter, Depends, Request, Response
 
-from jupyverse_api import Router
+from jupyverse_api import Config, Router
 
 from ..app import App
 from ..auth import Auth, User
@@ -61,3 +61,8 @@ def get_document(
         document_id: str,
     ):
         ...
+
+
+class YjsConfig(Config):
+    document_cleanup_delay: float = 60
+    document_save_delay: float = 1
diff --git a/plugins/yjs/fps_yjs/main.py b/plugins/yjs/fps_yjs/main.py
index eacd1b91..2028c0f8 100644
--- a/plugins/yjs/fps_yjs/main.py
+++ b/plugins/yjs/fps_yjs/main.py
@@ -8,22 +8,27 @@
 from jupyverse_api.app import App
 from jupyverse_api.auth import Auth
 from jupyverse_api.contents import Contents
-from jupyverse_api.yjs import Yjs
+from jupyverse_api.yjs import Yjs, YjsConfig
 
 from .routes import _Yjs
 
 
 class YjsComponent(Component):
+    def __init__(self, **kwargs):
+        self.yjs_config = YjsConfig(**kwargs)
+
     @context_teardown
     async def start(
         self,
         ctx: Context,
     ) -> AsyncGenerator[None, Optional[BaseException]]:
+        ctx.add_resource(self.yjs_config, types=YjsConfig)
+
         app = await ctx.request_resource(App)
         auth = await ctx.request_resource(Auth)  # type: ignore
         contents = await ctx.request_resource(Contents)  # type: ignore
 
-        yjs = _Yjs(app, auth, contents)
+        yjs = _Yjs(app, self.yjs_config, auth, contents)
         ctx.add_resource(yjs, types=Yjs)
 
         # start indexing in the background
diff --git a/plugins/yjs/fps_yjs/routes.py b/plugins/yjs/fps_yjs/routes.py
index 1a023d95..93fb000e 100644
--- a/plugins/yjs/fps_yjs/routes.py
+++ b/plugins/yjs/fps_yjs/routes.py
@@ -20,7 +20,7 @@
 from jupyverse_api.app import App
 from jupyverse_api.auth import Auth, User
 from jupyverse_api.contents import Contents
-from jupyverse_api.yjs import Yjs
+from jupyverse_api.yjs import Yjs, YjsConfig
 from jupyverse_api.yjs.models import CreateDocumentSession
 
 from .ydocs import ydocs as YDOCS
@@ -44,12 +44,13 @@ class _Yjs(Yjs):
     def __init__(
         self,
         app: App,
+        yjs_config: YjsConfig,
         auth: Auth,
         contents: Contents,
     ) -> None:
         super().__init__(app=app, auth=auth)
         self.contents = contents
-        self.room_manager = RoomManager(contents)
+        self.room_manager = RoomManager(contents, yjs_config)
         if Widgets is None:
             self.widgets = None
         else:
@@ -148,8 +149,9 @@ class RoomManager:
     websocket_server: JupyterWebsocketServer
     lock: asyncio.Lock
 
-    def __init__(self, contents: Contents):
+    def __init__(self, contents: Contents, yjs_config: YjsConfig):
         self.contents = contents
+        self.yjs_config = yjs_config
         self.documents = {}  # a dictionary of room_name:document
         self.watchers = {}  # a dictionary of file_id:task
         self.savers = {}  # a dictionary of file_id:task
@@ -314,7 +316,7 @@ async def maybe_save_document(
         self, file_id: str, file_type: str, file_format: str, document: YBaseDoc
     ) -> None:
         # save after 1 second of inactivity to prevent too frequent saving
-        await asyncio.sleep(1)  # FIXME: pass in config
+        await asyncio.sleep(self.yjs_config.document_save_delay)
         # if the room cannot be found, don't save
         try:
             file_path = await self.get_file_path(file_id, document)
@@ -349,12 +351,12 @@ async def maybe_save_document(
         del self.savers[file_id]
 
     async def maybe_clean_room(self, room, ws_path: str) -> None:
-        file_id = ws_path.split(":", 2)[2]
         # keep the document for a while in case someone reconnects
-        await asyncio.sleep(60)  # FIXME: pass in config
+        await asyncio.sleep(self.yjs_config.document_cleanup_delay)
         document = self.documents[ws_path]
         document.unobserve()
         del self.documents[ws_path]
+        file_id = ws_path.split(":", 2)[2]
         documents = [v for k, v in self.documents.items() if k.split(":", 2)[2] == file_id]
         if not documents:
             self.watchers[file_id].cancel()