feat: Add Document State Replacement for document-versioning based applications to y-partykit

.changeset/dry-pens-visit.md

@@ -0,0 +1,5 @@
+---
+"y-partyserver": patch
+---
+
+Add Document State Replacement for document-versioning based applications

packages/y-partyserver/src/server/index.ts

@@ -5,10 +5,28 @@ import type { Connection, ConnectionContext, WSMessage } from "partyserver";
 import { Server } from "partyserver";
 import * as awarenessProtocol from "y-protocols/awareness";
 import * as syncProtocol from "y-protocols/sync";
-import { applyUpdate, Doc as YDoc, encodeStateAsUpdate } from "yjs";
+import {
+  applyUpdate,
+  Doc as YDoc,
+  encodeStateAsUpdate,
+  encodeStateVector,
+  UndoManager,
+  XmlText,
+  XmlElement,
+  XmlFragment
+} from "yjs";
 
 import { handleChunked } from "../shared/chunking";
 
+const snapshotOrigin = Symbol("snapshot-origin");
+type YjsRootType =
+  | "Text"
+  | "Map"
+  | "Array"
+  | "XmlText"
+  | "XmlElement"
+  | "XmlFragment";
+
 const wsReadyStateConnecting = 0;
 const wsReadyStateOpen = 1;
 const wsReadyStateClosing = 2;
@@ -157,14 +175,78 @@ export class YServer<Env = unknown> extends Server<Env> {
   static callbackOptions: CallbackOptions = {};
 
   #ParentClass: typeof YServer = Object.getPrototypeOf(this).constructor;
-  readonly document = new WSSharedDoc();
+  readonly document: WSSharedDoc = new WSSharedDoc();
 
   async onLoad(): Promise<void> {
     // to be implemented by the user
     return;
   }
 
-  async onSave(): Promise<void> {}
+  async onSave(): Promise<void> {
+    // to be implemented by the user
+  }
+
+  /**
+   * Replaces the document with a different state using Yjs UndoManager key remapping.
+   *
+   * @param snapshotUpdate - The snapshot update to replace the document with.
+   * @param getMetadata (optional) - A function that returns the type of the root for a given key.
+   */
+  unstable_replaceDocument(
+    snapshotUpdate: Uint8Array,
+    getMetadata: (key: string) => YjsRootType = () => "Map"
+  ): void {
+    try {
+      const doc = this.document;
+      const snapshotDoc = new YDoc();
+      applyUpdate(snapshotDoc, snapshotUpdate, snapshotOrigin);
+
+      const currentStateVector = encodeStateVector(doc);
+      const snapshotStateVector = encodeStateVector(snapshotDoc);
+
+      const changesSinceSnapshotUpdate = encodeStateAsUpdate(
+        doc,
+        snapshotStateVector
+      );
+
+      const undoManager = new UndoManager(
+        [...snapshotDoc.share.keys()].map((key) => {
+          const type = getMetadata(key);
+          if (type === "Text") {
+            return snapshotDoc.getText(key);
+          } else if (type === "Map") {
+            return snapshotDoc.getMap(key);
+          } else if (type === "Array") {
+            return snapshotDoc.getArray(key);
+          } else if (type === "XmlText") {
+            return snapshotDoc.get(key, XmlText);
+          } else if (type === "XmlElement") {
+            return snapshotDoc.get(key, XmlElement);
+          } else if (type === "XmlFragment") {
+            return snapshotDoc.get(key, XmlFragment);
+          }
+          throw new Error(`Unknown root type: ${type} for key: ${key}`);
+        }),
+        {
+          trackedOrigins: new Set([snapshotOrigin])
+        }
+      );
+
+      applyUpdate(snapshotDoc, changesSinceSnapshotUpdate, snapshotOrigin);
+      undoManager.undo();
+
+      const documentChangesSinceSnapshotUpdate = encodeStateAsUpdate(
+        snapshotDoc,
+        currentStateVector
+      );
+
+      applyUpdate(this.document, documentChangesSinceSnapshotUpdate);
+    } catch (error) {
+      throw new Error(
+        `Failed to replace document: ${error instanceof Error ? error.message : "Unknown error"}`
+      );
+    }
+  }
 
   async onStart(): Promise<void> {
     const src = await this.onLoad();

packages/y-partyserver/src/shared/utils.ts

@@ -0,0 +1,30 @@
+// This file contains a shared implementation of base64 to uint8Array and uint8Array to base64.
+// Because certain text documents may be quite large, we split them into chunks of 8192 bytes to encode/decode.
+
+export function base64ToUint8Array(base64: string): Uint8Array {
+  const binaryString = atob(base64);
+  const uint8Array = new Uint8Array(binaryString.length);
+
+  const chunkSize = 8192;
+
+  for (let i = 0; i < binaryString.length; i += chunkSize) {
+    const end = Math.min(i + chunkSize, binaryString.length);
+    for (let j = i; j < end; j++) {
+      uint8Array[j] = binaryString.charCodeAt(j);
+    }
+  }
+
+  return uint8Array;
+}
+
+export function uint8ArrayToBase64(uint8Array: Uint8Array): string {
+  let binaryString = "";
+  const chunkSize = 8192;
+
+  for (let i = 0; i < uint8Array.length; i += chunkSize) {
+    const chunk = uint8Array.slice(i, i + chunkSize);
+    binaryString += String.fromCharCode.apply(null, Array.from(chunk));
+  }
+
+  return btoa(binaryString);
+}