Use jupyter-ydoc's aset_progressively

.github/workflows/test.yml

@@ -306,6 +306,7 @@ jobs:
         run: |
           jlpm test
           TIMELINE_FEATURE=1 jlpm test:timeline
+          CONFLICT_FEATURE=1 jlpm test:conflict
 
       - name: Upload Playwright Test report
         if: always()

.pre-commit-config.yaml

@@ -44,7 +44,7 @@ repos:
       - id: mypy
         exclude: "(^binder/jupyter_config\\.py$)|(^scripts/bump_version\\.py$)|(/setup\\.py$)"
         args: ["--config-file", "pyproject.toml"]
-        additional_dependencies: [tornado, pytest, 'pycrdt-websocket>=0.16.0,<0.17.0']
+        additional_dependencies: [tornado, pytest, 'pycrdt-websocket>=0.16.4,<0.17.0']
         stages: [manual]
 
   - repo: https://github.com/python-jsonschema/check-jsonschema

docs/source/configuration.md

@@ -52,6 +52,13 @@ jupyter lab --YDocExtension.file_poll_interval=2
 # If None, the document will be kept in memory forever.
 jupyter lab --YDocExtension.document_cleanup_delay=100
 
+# Whether documents loaded from disk should be streamed progressively (default: False).
+jupyter lab --YDocExtension.document_load_progressively=True
+
+# Delay notebook outputs larger than this size in MB during progressive loading (default: 100).
+# If None, outputs are loaded with the inputs.
+jupyter lab --YDocExtension.notebook_output_delay_threshold_mb=50
+
 # The YStore class to use for storing Y updates (default: JupyterSQLiteYStore).
 jupyter lab --YDocExtension.ystore_class=pycrdt.store.TempFileYStore
 

projects/jupyter-server-ydoc/jupyter_server_ydoc/app.py

@@ -76,6 +76,22 @@ class YDocExtension(ExtensionApp):
         Defaults to 1s, if None then the document will never be saved.""",
     )
 
+    document_load_progressively = Bool(
+        False,
+        config=True,
+        help="""Whether to progressively load documents from disk into the shared document.
+        When enabled, a shared document can stream its content to clients before the full
+        source file has finished loading.""",
+    )
+
+    notebook_output_delay_threshold_mb = Float(
+        100,
+        allow_none=True,
+        config=True,
+        help="""Output size in MB above which a shared notebook may delay loading outputs during
+        progressive document loading. Set to None to keep loading outputs with the inputs.""",
+    )
+
     ystore_class = Type(
         default_value=SQLiteYStore,
         klass=BaseYStore,
@@ -118,6 +134,10 @@ def initialize_settings(self):
                 "collaborative_file_poll_interval": self.file_poll_interval,
                 "collaborative_document_cleanup_delay": self.document_cleanup_delay,
                 "collaborative_document_save_delay": self.document_save_delay,
+                "collaborative_document_load_progressively": self.document_load_progressively,
+                "collaborative_notebook_output_delay_threshold_mb": (
+                    self.notebook_output_delay_threshold_mb
+                ),
                 "collaborative_ystore_class": self.ystore_class,
                 "collaborative_session_store_path": self.session_store_path,
             }
@@ -166,6 +186,10 @@ def initialize_handlers(self):
                     {
                         "document_cleanup_delay": self.document_cleanup_delay,
                         "document_save_delay": self.document_save_delay,
+                        "document_load_progressively": self.document_load_progressively,
+                        "notebook_output_delay_threshold_mb": (
+                            self.notebook_output_delay_threshold_mb
+                        ),
                         "file_loaders": self.file_loaders,
                         "ystore_class": ystore_class,
                         "ywebsocket_server": self.ywebsocket_server,
@@ -258,11 +282,13 @@ async def get_document(
                     self.log,
                     exception_handler=exception_logger,
                     save_delay=self.document_save_delay,
+                    document_load_progressively=self.document_load_progressively,
+                    notebook_output_delay_threshold_mb=self.notebook_output_delay_threshold_mb,
                 )
-                await room.initialize()
                 try:
                     await self.ywebsocket_server.start_room(room)
                     self.ywebsocket_server.add_room(room_id, room)
+                    await room.initialize()
                     self.log.info(f"Created and started room: {room_id}")
                 except Exception as e:
                     self.log.error("Room %s failed to start on websocket server", room_id)

projects/jupyter-server-ydoc/jupyter_server_ydoc/handlers.py

@@ -137,6 +137,10 @@ def exception_logger(exception: Exception, log: Logger) -> bool:
                         self.log,
                         exception_handler=exception_logger,
                         save_delay=self._document_save_delay,
+                        document_load_progressively=self._document_load_progressively,
+                        notebook_output_delay_threshold_mb=(
+                            self._notebook_output_delay_threshold_mb
+                        ),
                     )
 
                 else:
@@ -181,6 +185,8 @@ def initialize(
         room_locks: dict[str, asyncio.Lock] | None = None,
         document_cleanup_delay: float | None = 60.0,
         document_save_delay: float | None = 1.0,
+        document_load_progressively: bool = False,
+        notebook_output_delay_threshold_mb: float | None = 100,
     ) -> None:
         self._background_tasks = set()
         # File ID manager cannot be passed as argument as the extension may load after this one
@@ -189,6 +195,8 @@ def initialize(
         self._ystore_class = ystore_class
         self._cleanup_delay = document_cleanup_delay
         self._document_save_delay = document_save_delay
+        self._document_load_progressively = document_load_progressively
+        self._notebook_output_delay_threshold_mb = notebook_output_delay_threshold_mb
         self._websocket_server = ywebsocket_server
         self._message_queue = asyncio.Queue()
         self._room_id = ""

projects/jupyter-server-ydoc/jupyter_server_ydoc/pytest_plugin.py

@@ -317,6 +317,7 @@ def _inner(
                 store,
                 None,
                 save_delay,
+                document_load_progressively=False,
             ),
         )
 

projects/jupyter-server-ydoc/jupyter_server_ydoc/rooms.py

@@ -40,6 +40,8 @@ def __init__(
         ystore: BaseYStore | None,
         log: Logger | None,
         save_delay: float | None = None,
+        document_load_progressively: bool = False,
+        notebook_output_delay_threshold_mb: float | None = 100,
         exception_handler: Callable[[Exception, Logger], bool] | None = None,
     ):
         super().__init__(ready=False, ystore=ystore, exception_handler=exception_handler, log=log)
@@ -53,6 +55,8 @@ def __init__(
 
         self._logger = logger
         self._save_delay = save_delay
+        self._document_load_progressively = document_load_progressively
+        self._notebook_output_delay_threshold_mb = notebook_output_delay_threshold_mb
 
         self._update_lock = asyncio.Lock()
         self._cleaner: asyncio.Task | None = None
@@ -119,7 +123,9 @@ async def initialize(self) -> None:
 
         model = await self._file.load_content(self._file_format, self._file_type)
 
-        async with self._update_lock:
+        await self._update_lock.acquire()
+        release_update_lock = True
+        try:
             # try to apply Y updates from the YStore for this document
             read_from_source = True
             loaded_from_store = False
@@ -171,18 +177,58 @@ async def initialize(self) -> None:
                     self._file.path,
                 )
                 if not loaded_from_store:
-                    await self._apply_deterministic_source_content(model["content"])
+                    if self._document_load_progressively:
+                        release_update_lock = False
+                        initialized = asyncio.Event()
+                        finish = asyncio.Event()
+                        self.create_task(
+                            self._finish_progressive_initialization(
+                                model["content"], initialized, finish
+                            )
+                        )
+                        await initialized.wait()
+                        self.ready = True
+                        await self.ydoc_observed.wait()
+                        finish.set()
+                        self._emit(LogLevel.INFO, "initialize", "Room initialized")
+                        return
+                    else:
+                        await self._apply_deterministic_source_content(model["content"])
                 else:
                     await self._document.aset(model["content"])
 
                 if self.ystore:
                     await self.ystore.encode_state_as_update(self.ydoc)
 
-            self._document.dirty = False
             self.ready = True
             self._emit(LogLevel.INFO, "initialize", "Room initialized")
+        finally:
+            if release_update_lock:
+                self._update_lock.release()
 
-    async def _apply_deterministic_source_content(self, content: Any) -> None:
+    async def _finish_progressive_initialization(
+        self, content: Any, initialized: asyncio.Event, finish: asyncio.Event
+    ) -> None:
+        try:
+            await self._apply_deterministic_source_content(
+                content, progressive=True, initialized=initialized, finish=finish
+            )
+            if self.ystore:
+                await self.ystore.encode_state_as_update(self.ydoc)
+        except Exception as e:
+            msg = f"Error loading content from file: {self._file.path}\n{e!r}"
+            self.log.error(msg, exc_info=e)
+            self._emit(LogLevel.ERROR, None, msg)
+        finally:
+            self._update_lock.release()
+
+    async def _apply_deterministic_source_content(
+        self,
+        content: Any,
+        progressive: bool = False,
+        initialized: asyncio.Event | None = None,
+        finish: asyncio.Event | None = None,
+    ) -> None:
         """Load source content using a deterministic update.
 
         Rooms rebuilt from disk must recreate the same Yjs history for identical
@@ -194,8 +240,20 @@ async def _apply_deterministic_source_content(self, content: Any) -> None:
         """
         source_ydoc: Doc = Doc(client_id=0)
         source_document = YDOCS.get(self._file_type, YFILE)(source_ydoc)
-        await source_document.aset(content)
-        self.ydoc.apply_update(source_ydoc.get_update())
+        if progressive:
+            subscription = source_ydoc.observe(lambda event: self.ydoc.apply_update(event.update))
+            try:
+                await source_document.aset_progressively(
+                    content,
+                    initialized=initialized,
+                    finish=finish,
+                    delay_outputs_above_mb=self._notebook_output_delay_threshold_mb,
+                )
+            finally:
+                source_ydoc.unobserve(subscription)
+        else:
+            await source_document.aset(content)
+            self.ydoc.apply_update(source_ydoc.get_update())
 
     def _emit(self, level: LogLevel, action: str | None = None, msg: str | None = None) -> None:
         data = {"level": level.value, "room": self._room_id, "path": self._file.path}

projects/jupyter-server-ydoc/pyproject.toml

@@ -29,9 +29,9 @@ authors = [
 ]
 dependencies = [
     "jupyter_server>=2.15.0,<3.0.0",
-    "jupyter_ydoc>=4.0.0b0,<5.0.0",
+    "jupyter_ydoc >=4.1.1,<5.0.0",
     "pycrdt",
-    "pycrdt-websocket>=0.16.2,<0.17.0",
+    "pycrdt-websocket>=0.16.4,<0.17.0",
     "jupyter_events>=0.11.0",
     "jupyter_server_fileid>=0.7.0,<1",
     "jsonschema>=4.18.0"
@@ -71,3 +71,6 @@ exclude = ["/.github", "/binder", "node_modules"]
 
 [tool.check-wheel-contents]
 ignore = ["W002"]
+
+[tool.hatch.metadata]
+allow-direct-references = true

tests/test_app.py

@@ -16,6 +16,8 @@ def test_default_settings(jp_serverapp):
     assert settings["file_poll_interval"] == 1
     assert settings["document_cleanup_delay"] == 60
     assert settings["document_save_delay"] == 1
+    assert settings["document_load_progressively"] is False
+    assert settings["notebook_output_delay_threshold_mb"] == 100
     assert settings["ystore_class"] == SQLiteYStore
 
 
@@ -55,6 +57,24 @@ def test_settings_should_change_save_delay(jp_configurable_serverapp):
     assert settings["document_save_delay"] == 10
 
 
+def test_settings_should_change_document_progressive_loading(jp_configurable_serverapp):
+    argv = ["--YDocExtension.document_load_progressively=True"]
+
+    app = jp_configurable_serverapp(argv=argv)
+    settings = app.web_app.settings["jupyter_server_ydoc_config"]
+
+    assert settings["document_load_progressively"] is True
+
+
+def test_settings_should_change_notebook_output_delay_threshold(jp_configurable_serverapp):
+    argv = ["--YDocExtension.notebook_output_delay_threshold_mb=50"]
+
+    app = jp_configurable_serverapp(argv=argv)
+    settings = app.web_app.settings["jupyter_server_ydoc_config"]
+
+    assert settings["notebook_output_delay_threshold_mb"] == 50
+
+
 def test_settings_should_change_ystore_class(jp_configurable_serverapp):
     argv = ["--YDocExtension.ystore_class=jupyter_server_ydoc.stores.TempFileYStore"]
 

tests/test_documents.py

@@ -64,9 +64,14 @@ async def test_dirty(
 
     websocket, room_name = await rtc_connect_doc_client(file_format, file_type, file_path)
     async with websocket as ws, Provider(jupyter_ydoc.ydoc, HttpxWebsocket(ws, room_name)):
-        for _ in range(2):
+        for _ in range(3):
             jupyter_ydoc.dirty = True
             await sleep(rtc_document_save_delay * 1.5)
+            # the server might not be ready to receive updates from this client yet
+            # if it is, it should clear the dirty state
+            if not jupyter_ydoc.dirty:
+                return
+        else:
             assert not jupyter_ydoc.dirty
 
 
@@ -164,6 +169,7 @@ async def _create_notebook_room(notebook: dict, room_id: str) -> tuple[DocumentR
         None,
         None,
         None,
+        document_load_progressively=False,
     )
     await room.initialize()
     return room, loader

ui-tests/jupyter_server_test_config.py

@@ -19,6 +19,11 @@
 # Fast room eviction so conflict tests don't need to wait 60 seconds.
 c.YDocExtension.document_cleanup_delay = 1
 
+# Keep the delayed-output path observable in UI tests without oversized
+# notebook fixtures.
+c.YDocExtension.document_load_progressively = True
+c.YDocExtension.notebook_output_delay_threshold_mb = 7
+
 # Force-close dead WebSocket connections quickly.  Playwright's setOffline(true)
 # blocks network I/O without tearing down existing TCP connections, so pings are
 # needed to make the server detect the disconnection.  The conflict test goes

ui-tests/jupyter_server_test_config_noprogressive.py

@@ -0,0 +1,44 @@
+# Copyright (c) Jupyter Development Team.
+# Distributed under the terms of the Modified BSD License.
+
+"""Server configuration for conflict-resolution integration tests.
+
+Progressive document loading is disabled so that
+_apply_deterministic_source_content fully populates the Ydoc before
+any client sync, which is needed for the stale-parent RuntimeError
+("block parent") to be raised on reconnect after an out-of-band change.
+
+!! Never use this configuration in production because it
+opens the server to the world and provide access to JupyterLab
+JavaScript objects through the global window variable.
+"""
+
+from typing import Any, cast
+
+from jupyterlab.galata import configure_jupyter_server
+
+# `c` is injected by jupyter-server while loading this config file.
+c = cast(Any, globals()["c"])
+configure_jupyter_server(c)  # noqa
+
+# Fast room eviction so conflict tests don't need to wait 60 seconds.
+c.YDocExtension.document_cleanup_delay = 1
+
+# Disable progressive loading – the full document content must be
+# deterministically loaded before the client can sync so that stale
+# parent references trigger the conflict detection path.
+c.YDocExtension.document_load_progressively = False
+
+# Force-close dead WebSocket connections quickly.  Playwright's setOffline(true)
+# blocks network I/O without tearing down existing TCP connections, so pings are
+# needed to make the server detect the disconnection.  The conflict test goes
+# offline for 10 s; with interval=2 + timeout=5 the dead connection closes ≤7 s
+# after going offline, leaving enough margin before the test comes back online.
+c.ServerApp.websocket_ping_interval = 2  # seconds between pings
+c.ServerApp.websocket_ping_timeout = 5  # close connection if no pong within 5 s
+
+# Use SQLiteYStore with a predictable path so conflict tests can delete
+# the database during the offline period to force room recreation via
+# _apply_deterministic_source_content. This simulates the production scenario
+# where a room is evicted and the notebook structure changes on disk.
+c.SQLiteYStore.db_path = "/tmp/jupyter_ystore_ui_test.db"

ui-tests/package.json

@@ -5,8 +5,10 @@
   "private": true,
   "scripts": {
     "start": "jupyter lab --config jupyter_server_test_config.py",
+    "start:conflict": "jupyter lab --config jupyter_server_test_config_noprogressive.py",
     "start:timeline": "jupyter lab --config jupyter_server_test_config.py --ServerApp.base_url=/api/collaboration/timeline/",
     "test": "npx playwright test --config=playwright.config.js",
+    "test:conflict": "npx playwright test --config=playwright.conflict.config.js",
     "test:timeline": "npx playwright test --config=playwright.timeline.config.js",
     "test:update": "npx playwright test --update-snapshots"
   },