diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index ecc138ce..1024fdd3 100644 --- a/.github/workflows/test.yml +++ b/.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() diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 7b4fa2ea..207cdf66 100644 --- a/.pre-commit-config.yaml +++ b/.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 diff --git a/docs/source/configuration.md b/docs/source/configuration.md index 0f7e8115..79edaf95 100644 --- a/docs/source/configuration.md +++ b/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 diff --git a/projects/jupyter-server-ydoc/jupyter_server_ydoc/app.py b/projects/jupyter-server-ydoc/jupyter_server_ydoc/app.py index ba25a12e..da2ee4a9 100644 --- a/projects/jupyter-server-ydoc/jupyter_server_ydoc/app.py +++ b/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) diff --git a/projects/jupyter-server-ydoc/jupyter_server_ydoc/handlers.py b/projects/jupyter-server-ydoc/jupyter_server_ydoc/handlers.py index a63bd379..588e0646 100644 --- a/projects/jupyter-server-ydoc/jupyter_server_ydoc/handlers.py +++ b/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 = "" diff --git a/projects/jupyter-server-ydoc/jupyter_server_ydoc/pytest_plugin.py b/projects/jupyter-server-ydoc/jupyter_server_ydoc/pytest_plugin.py index dec1d1be..a1c28df8 100644 --- a/projects/jupyter-server-ydoc/jupyter_server_ydoc/pytest_plugin.py +++ b/projects/jupyter-server-ydoc/jupyter_server_ydoc/pytest_plugin.py @@ -317,6 +317,7 @@ def _inner( store, None, save_delay, + document_load_progressively=False, ), ) diff --git a/projects/jupyter-server-ydoc/jupyter_server_ydoc/rooms.py b/projects/jupyter-server-ydoc/jupyter_server_ydoc/rooms.py index d432933e..ad71e65b 100644 --- a/projects/jupyter-server-ydoc/jupyter_server_ydoc/rooms.py +++ b/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} diff --git a/projects/jupyter-server-ydoc/pyproject.toml b/projects/jupyter-server-ydoc/pyproject.toml index 8029d1c9..6395a692 100644 --- a/projects/jupyter-server-ydoc/pyproject.toml +++ b/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 diff --git a/tests/test_app.py b/tests/test_app.py index db9a6ffd..07a4fe2d 100644 --- a/tests/test_app.py +++ b/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"] diff --git a/tests/test_documents.py b/tests/test_documents.py index 50a285e9..4b6e701d 100644 --- a/tests/test_documents.py +++ b/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 diff --git a/ui-tests/jupyter_server_test_config.py b/ui-tests/jupyter_server_test_config.py index c6562035..d1bcad76 100644 --- a/ui-tests/jupyter_server_test_config.py +++ b/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 diff --git a/ui-tests/jupyter_server_test_config_noprogressive.py b/ui-tests/jupyter_server_test_config_noprogressive.py new file mode 100644 index 00000000..b508d225 --- /dev/null +++ b/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" diff --git a/ui-tests/package.json b/ui-tests/package.json index 59d97de3..ff94cfb0 100644 --- a/ui-tests/package.json +++ b/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" }, diff --git a/ui-tests/playwright.conflict.config.js b/ui-tests/playwright.conflict.config.js new file mode 100644 index 00000000..3815ccdb --- /dev/null +++ b/ui-tests/playwright.conflict.config.js @@ -0,0 +1,28 @@ +/* + * Copyright (c) Jupyter Development Team. + * Distributed under the terms of the Modified BSD License. + */ + +/** + * Configuration for Playwright using default from @jupyterlab/galata + */ +const baseConfig = require('@jupyterlab/galata/lib/playwright-config'); + +module.exports = { + ...baseConfig, + workers: 1, + testMatch: 'tests/conflict.spec.ts', + testIgnore: '**/.ipynb_checkpoints/**', + timeout: 120 * 1000, + webServer: { + command: 'jlpm start:conflict', + url: 'http://localhost:8888/lab', + timeout: 120 * 1000, + reuseExistingServer: !process.env.CI + }, + expect: { + toMatchSnapshot: { + maxDiffPixelRatio: 0.01 + } + } +}; diff --git a/ui-tests/tests/conflict.spec.ts b/ui-tests/tests/conflict.spec.ts index 210140f3..1cd15424 100644 --- a/ui-tests/tests/conflict.spec.ts +++ b/ui-tests/tests/conflict.spec.ts @@ -5,6 +5,9 @@ import { expect, galata, test } from '@jupyterlab/galata'; import { unlink } from 'fs/promises'; import type { Page, APIRequestContext } from '@playwright/test'; +const isConflictEnv = process.env.CONFLICT_FEATURE || "0"; +const isConflict = parseInt(isConflictEnv) + /** * A minimal notebook with one cell. * @@ -158,6 +161,10 @@ test.describe.serial('Conflict handling', () => { test( 'shows a conflict dialog and dismisses it', async ({ page, request, tmpPath, baseURL }) => { + if (!isConflict) { + console.log('Skipping this test.'); + return; + } const dialog = await triggerConflict( page, request, @@ -174,6 +181,10 @@ test.describe.serial('Conflict handling', () => { test( 'Revert button reloads the document from disk', async ({ page, request, tmpPath, baseURL }) => { + if (!isConflict) { + console.log('Skipping this test.'); + return; + } const dialog = await triggerConflict( page, request, @@ -207,6 +218,10 @@ test.describe.serial('Conflict handling', () => { test( 'Save As button opens the save-as dialog', async ({ page, request, tmpPath, baseURL }) => { + if (!isConflict) { + console.log('Skipping this test.'); + return; + } const dialog = await triggerConflict( page, request, @@ -229,6 +244,10 @@ test.describe.serial('Conflict handling', () => { test( 'Show Diff button opens a diff widget', async ({ page, request, tmpPath, baseURL }) => { + if (!isConflict) { + console.log('Skipping this test.'); + return; + } const dialog = await triggerConflict( page, request, @@ -249,6 +268,10 @@ test.describe.serial('Conflict handling', () => { test( 'Save Local As button in diff toolbar opens the save-as dialog', async ({ page, request, tmpPath, baseURL }) => { + if (!isConflict) { + console.log('Skipping this test.'); + return; + } const dialog = await triggerConflict( page, request, @@ -276,6 +299,10 @@ test.describe.serial('Conflict handling', () => { test( 'Revert to Remote button in diff toolbar reloads the document', async ({ page, request, tmpPath, baseURL }) => { + if (!isConflict) { + console.log('Skipping this test.'); + return; + } const dialog = await triggerConflict( page, request, diff --git a/ui-tests/tests/progressive-loading.spec.ts b/ui-tests/tests/progressive-loading.spec.ts new file mode 100644 index 00000000..25bda0e7 --- /dev/null +++ b/ui-tests/tests/progressive-loading.spec.ts @@ -0,0 +1,229 @@ +// Copyright (c) Jupyter Development Team. +// Distributed under the terms of the Modified BSD License. + +import { expect, galata, test } from '@jupyterlab/galata'; +import type { IJupyterLabPageFixture } from '@jupyterlab/galata'; +import type * as nbformat from '@jupyterlab/nbformat'; +import type { BrowserContext } from '@playwright/test'; +import { rm } from 'fs/promises'; + +const NOTEBOOK_NAME = 'progressive_loading.ipynb'; +const YSTORE_DB = '/tmp/jupyter_ystore_ui_test.db'; +const YSTORE_FILES = [YSTORE_DB, `${YSTORE_DB}-shm`, `${YSTORE_DB}-wal`]; +const HEAVY_PAYLOAD_BYTES = 8 * 1024 * 1024; +const HEAVY_PAYLOAD_ALPHABET = + 'ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789'; +const LARGE_COLLABORATION_MESSAGE_BYTES = 512 * 1024; +const DELAYED_OUTPUT_MARKERS = [ + 'delayed-output-1-ready', + 'delayed-output-2-ready', + 'delayed-output-3-ready' +]; + +function makeHeavyPayload(seed: number): string { + const chunks: string[] = []; + let state = seed; + + // Repeated text compresses too well over WebSocket to exercise the large + // message path reliably. + for (let i = 0; i < HEAVY_PAYLOAD_BYTES; i += 1024) { + let chunk = ''; + for (let j = 0; j < 1024; j++) { + state = (state * 1664525 + 1013904223) >>> 0; + chunk += HEAVY_PAYLOAD_ALPHABET[state % HEAVY_PAYLOAD_ALPHABET.length]; + } + chunks.push(chunk); + } + + return chunks.join(''); +} + +function heavyOutput(label: string, seed: number): nbformat.IDisplayData { + return { + output_type: 'display_data', + data: { + // Keep the payload renderable but invisible in the UI. + 'text/html': `
${label}
` + }, + metadata: {} + }; +} + +function collaborationMessageSize(message: string | Buffer): number { + return typeof message === 'string' + ? Buffer.byteLength(message) + : message.byteLength; +} + +const LARGE_OUTPUT_NOTEBOOK = galata.Notebook.makeNotebook([ + galata.Notebook.makeCell({ + cell_type: 'markdown', + id: 'intro', + source: '# Progressive loading smoke' + }), + galata.Notebook.makeCell({ + cell_type: 'code', + id: 'small-output', + source: 'print("small output")', + outputs: [ + { + output_type: 'stream', + name: 'stdout', + text: 'small-output-ready\n' + } + ], + execution_count: 1 + }), + galata.Notebook.makeCell({ + cell_type: 'code', + id: 'heavy-output-1', + source: 'display heavy output 1', + outputs: [heavyOutput(DELAYED_OUTPUT_MARKERS[0], 1)], + execution_count: 1 + }), + galata.Notebook.makeCell({ + cell_type: 'code', + id: 'heavy-output-2', + source: 'display heavy output 2', + outputs: [heavyOutput(DELAYED_OUTPUT_MARKERS[1], 2)], + execution_count: 1 + }), + galata.Notebook.makeCell({ + cell_type: 'code', + id: 'heavy-output-3', + source: 'display heavy output 3', + outputs: [heavyOutput(DELAYED_OUTPUT_MARKERS[2], 3)], + execution_count: 1 + }) +]); + +async function holdLargeCollaborationMessages(context: BrowserContext): Promise<{ + release: () => void; + heldMessageCount: () => number; +}> { + let released = false; + let heldMessages = 0; + const queuedMessages: Array<() => void> = []; + + await context.routeWebSocket( + url => url.pathname.includes('/api/collaboration/room/'), + ws => { + const server = ws.connectToServer(); + + ws.onMessage(message => { + server.send(message); + }); + server.onMessage(message => { + const forwardMessage = () => { + ws.send(message); + }; + + if ( + collaborationMessageSize(message) > LARGE_COLLABORATION_MESSAGE_BYTES + ) { + // Make the first-pass notebook state observable before outputs arrive. + if (released) { + forwardMessage(); + } else { + heldMessages++; + queuedMessages.push(forwardMessage); + } + } else { + forwardMessage(); + } + }); + } + ); + + return { + release: () => { + released = true; + for (const forwardMessage of queuedMessages.splice(0)) { + forwardMessage(); + } + }, + heldMessageCount: () => heldMessages + }; +} + +test.describe('Progressive notebook loading', () => { + test.beforeEach(async () => { + // The progressive path is skipped when the document is restored from ystore. + await Promise.all(YSTORE_FILES.map(file => rm(file, { force: true }))); + }); + + test.afterEach(async ({ request, tmpPath }) => { + const contents = galata.newContentsHelper(request); + await contents.deleteFile(`${tmpPath}/${NOTEBOOK_NAME}`).catch(() => {}); + }); + + test('renders notebook content before all delayed large outputs finish', async ({ + browser, + request, + tmpPath, + baseURL, + waitForApplication + }) => { + // Install the WebSocket route before creating the page; Galata's default + // page fixture is already initialized by the time the test body runs. + const context = await browser.newContext(); + const largeMessages = await holdLargeCollaborationMessages(context); + let page: IJupyterLabPageFixture | undefined; + + try { + page = await galata.initTestPage( + '/lab', + true, + baseURL!, + true, + galata.DEFAULT_SETTINGS, + true, + true, + await context.newPage(), + new Map(), + new Map(), + tmpPath, + waitForApplication, + new Map(), + true + ); + + // uploadContent creates a generic file; the progressive path needs an + // actual notebook model. + const response = await request.put( + `${baseURL}/api/contents/${tmpPath}/${NOTEBOOK_NAME}`, + { + headers: { 'Content-Type': 'application/json' }, + data: JSON.stringify({ + type: 'notebook', + format: 'json', + content: LARGE_OUTPUT_NOTEBOOK + }) + } + ); + expect(response.ok()).toBeTruthy(); + + await page.filebrowser.refresh(); + expect(await page.notebook.open(NOTEBOOK_NAME)).toBeTruthy(); + await expect(page.locator('.jp-Dialog:visible')).toHaveCount(0); + + const notebook = page.locator('.jp-Notebook'); + await expect(page.locator('.jp-Notebook .jp-Cell')).toHaveCount(5, { + timeout: 15000 + }); + await expect(notebook).toContainText('small-output-ready'); + expect(await notebook.textContent()).not.toContain( + DELAYED_OUTPUT_MARKERS[2] + ); + await expect.poll(largeMessages.heldMessageCount).toBeGreaterThan(0); + + largeMessages.release(); + for (const marker of DELAYED_OUTPUT_MARKERS) { + await expect(notebook).toContainText(marker, { timeout: 30000 }); + } + } finally { + await page?.close().catch(() => {}); + await context.close().catch(() => {}); + } + }); +});