Fix confirmation on manual saving (#490)

* Fix confirmation on manual saving

* Please `mypy`

* Use `MessageType.RAW` instead of a custom var

* Update error message

Co-authored-by: David Brochart <david.brochart@gmail.com>

* Pin pre-commit `pycrdt-websocket` version

* Add save ID to avoid mixing statuses

---------

Co-authored-by: David Brochart <david.brochart@gmail.com>

Author: krassowski

.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]
+        additional_dependencies: [tornado, pytest, 'pycrdt-websocket<0.16.0']
         stages: [manual]
 
   - repo: https://github.com/sirosen/check-jsonschema

packages/docprovider/src/ydrive.ts

@@ -11,6 +11,7 @@ import {
   ServerConnection,
   User
 } from '@jupyterlab/services';
+import { PromiseDelegate } from '@lumino/coreutils';
 import { ISignal, Signal } from '@lumino/signaling';
 
 import { DocumentChange, ISharedDocument, YDocument } from '@jupyter/ydoc';
@@ -22,20 +23,14 @@ import {
 } from '@jupyter/collaborative-drive';
 import { Awareness } from 'y-protocols/awareness';
 import { ISettingRegistry } from '@jupyterlab/settingregistry';
+import * as decoding from 'lib0/decoding';
 import * as encoding from 'lib0/encoding';
 
 const DISABLE_RTC =
   PageConfig.getOption('disableRTC') === 'true' ? true : false;
 
 const RAW_MESSAGE_TYPE = 2;
 
-const SAVE_MESSAGE = (() => {
-  const encoder = encoding.createEncoder();
-  encoding.writeVarUint(encoder, RAW_MESSAGE_TYPE);
-  encoding.writeVarString(encoder, 'save');
-  return encoding.toUint8Array(encoder);
-})();
-
 /**
  * The url for the default drive service.
  */
@@ -134,9 +129,60 @@ export class RtcContentProvider
     if (options.format && options.type) {
       const key = `${options.format}:${options.type}:${localPath}`;
       const provider = this._providers.get(key);
+      const saveId = ++this._saveCounter;
 
       if (provider) {
-        provider.wsProvider?.ws?.send(SAVE_MESSAGE);
+        const ws = provider.wsProvider?.ws;
+        if (ws) {
+          const delegate = new PromiseDelegate<void>();
+          const handler = (event: MessageEvent) => {
+            const data = new Uint8Array(event.data);
+            const decoder = decoding.createDecoder(data);
+            try {
+              const messageType = decoding.readVarUint(decoder);
+              if (messageType !== RAW_MESSAGE_TYPE) {
+                return;
+              }
+            } catch {
+              return;
+            }
+            const rawReply = decoding.readVarString(decoder);
+            let reply: {
+              type: 'save';
+              responseTo: number;
+              status: 'success' | 'skipped' | 'failed';
+            } | null = null;
+            try {
+              reply = JSON.parse(rawReply);
+            } catch (e) {
+              console.debug('The raw reply received was not a JSON reply');
+            }
+            if (
+              reply &&
+              reply['type'] === 'save' &&
+              reply['responseTo'] === saveId
+            ) {
+              if (reply.status === 'success') {
+                delegate.resolve();
+              } else if (reply.status === 'failed') {
+                delegate.reject('Saving failed');
+              } else if (reply.status === 'skipped') {
+                delegate.reject('Saving already in progress');
+              } else {
+                delegate.reject('Unrecognised save reply status');
+              }
+            }
+          };
+          ws.addEventListener('message', handler);
+          const encoder = encoding.createEncoder();
+          encoding.writeVarUint(encoder, RAW_MESSAGE_TYPE);
+          encoding.writeVarString(encoder, 'save');
+          encoding.writeVarUint(encoder, saveId);
+          const saveMessage = encoding.toUint8Array(encoder);
+          ws.send(saveMessage);
+          await delegate.promise;
+          ws.removeEventListener('message', handler);
+        }
         const fetchOptions: Contents.IFetchOptions = {
           type: options.type,
           format: options.format,
@@ -256,6 +302,7 @@ export class RtcContentProvider
   };
 
   private _user: User.IManager;
+  private _saveCounter = 0;
   private _trans: TranslationBundle;
   private _globalAwareness: Awareness | null;
   private _providers: Map<string, WebSocketProvider>;

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

@@ -15,7 +15,7 @@
 from jupyter_server.base.handlers import APIHandler, JupyterHandler
 from jupyter_server.utils import ensure_async
 from jupyter_ydoc import ydocs as YDOCS
-from pycrdt import Doc, UndoManager
+from pycrdt import Doc, Encoder, UndoManager
 from pycrdt_websocket.yroom import YRoom
 from pycrdt_websocket.ystore import BaseYStore
 from tornado import web
@@ -273,7 +273,7 @@ async def open(self, room_id: str) -> None:  # type:ignore[override]
             if self._room_id != "JupyterLab:globalAwareness":
                 self._emit_awareness_event(self.current_user.username, "join")
 
-    async def send(self, message):
+    async def send(self, message: bytes) -> None:
         """
         Send a message to the client.
         """
@@ -299,16 +299,38 @@ async def on_message(self, message):
         if header == MessageType.RAW:
             msg = decoder.read_var_string()
             if msg == "save":
+                save_id = decoder.read_var_uint()
+                save_reply = {
+                    "type": "save",
+                    "responseTo": save_id,
+                }
                 try:
                     room = cast(DocumentRoom, self.room)
-                    room._save_to_disc()
+                    save_task = room._save_to_disc()
+                    # task may be missing if save was already in progress
+                    if save_task:
+                        await save_task
+                        await self.send(
+                            self._encode_json_message({**save_reply, "status": "success"})
+                        )
+                    else:
+                        await self.send(
+                            self._encode_json_message({**save_reply, "status": "skipped"})
+                        )
                 except Exception:
                     self.log.error("Couldn't save content from room: %s", self._room_id)
+                    await self.send(self._encode_json_message({**save_reply, "status": "failed"}))
             return
 
         self._message_queue.put_nowait(message)
         self._websocket_server.ypatch_nb += 1
 
+    def _encode_json_message(self, message: dict) -> bytes:
+        encoder = Encoder()
+        encoder.write_var_uint(MessageType.RAW)
+        encoder.write_var_string(json.dumps(message))
+        return encoder.to_bytes()
+
     def on_close(self) -> None:
         """
         On connection close.

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

@@ -280,6 +280,7 @@ def _save_to_disc(self):
         self._saving_document = asyncio.create_task(
             self._maybe_save_document(self._saving_document)
         )
+        return self._saving_document
 
     async def _maybe_save_document(self, saving_document: asyncio.Task | None) -> None:
         """