fix(core): avoid torn Yjs plugin state during fork merge

packages/core/src/editor/managers/ExtensionManager/index.ts

@@ -7,7 +7,7 @@ import {
   Extension as TiptapExtension,
 } from "@tiptap/core";
 import { keydownHandler } from "@tiptap/pm/keymap";
-import { Plugin, TextSelection } from "prosemirror-state";
+import { Plugin, PluginKey, TextSelection } from "prosemirror-state";
 import { updateBlockTr } from "../../../api/blockManipulation/commands/updateBlock/updateBlock.js";
 import { setTextCursorPosition } from "../../../api/blockManipulation/selections/textCursorPosition.js";
 import {
@@ -260,6 +260,7 @@ export class ExtensionManager {
    * Atomically replace extension instances in the editor.
    * @param toUnregister - The extensions to unregister, can be a string key, an extension instance, an extension factory, or an array of any of those
    * @param toRegister - The extensions to register, can be an extension instance, an extension factory, or an array of any of those
+   * @param options.resetPluginStateFor - Plugin keys whose state must be initialized afresh instead of carried across the replacement
    * @returns void
    */
   public replaceExtension(
@@ -273,6 +274,7 @@ export class ExtensionManager {
       | Extension
       | ExtensionFactoryInstance
       | (Extension | ExtensionFactoryInstance)[],
+    options?: { resetPluginStateFor: readonly PluginKey[] },
   ): void {
     // ---- Remove phase (no updatePlugins call) ----
     const extensionsToRemove = this.resolveExtensions(toUnregister);
@@ -365,36 +367,53 @@ export class ExtensionManager {
     }
 
     // ---- Single atomic plugin update ----
-    this.updatePlugins((plugins) => [
-      ...plugins.filter((plugin) => {
-        // Fast path: exact reference match
-        if (pluginRefsToRemove.has(plugin)) {
-          return false;
-        }
-        // Fallback: match by key string (handles cases where plugin instances
-        // in the state differ from the ones we tracked)
-        if (pluginKeysToRemove.size) {
-          const key = (plugin as any).spec?.key;
-          const keyStr = typeof key === "object" && key ? key.key : key;
-          if (typeof keyStr === "string" && pluginKeysToRemove.has(keyStr)) {
+    this.updatePlugins(
+      (plugins) => [
+        ...plugins.filter((plugin) => {
+          // Fast path: exact reference match
+          if (pluginRefsToRemove.has(plugin)) {
             return false;
           }
-        }
-        return true;
-      }),
-      ...pluginsToAdd,
-    ]);
+          // Fallback: match by key string (handles cases where plugin instances
+          // in the state differ from the ones we tracked)
+          if (pluginKeysToRemove.size) {
+            const key = (plugin as any).spec?.key;
+            const keyStr = typeof key === "object" && key ? key.key : key;
+            if (typeof keyStr === "string" && pluginKeysToRemove.has(keyStr)) {
+              return false;
+            }
+          }
+          return true;
+        }),
+        ...pluginsToAdd,
+      ],
+      options?.resetPluginStateFor,
+    );
   }
 
   /**
    * Allows resetting the current prosemirror state's plugins
    * @param update - A function that takes the current plugins and returns the new plugins
    * @returns void
    */
-  private updatePlugins(update: (plugins: Plugin[]) => Plugin[]): void {
+  private updatePlugins(
+    update: (plugins: Plugin[]) => Plugin[],
+    resetPluginStateFor?: readonly PluginKey[],
+  ): void {
     const currentState = this.editor.prosemirrorState;
-
-    const state = currentState.reconfigure({
+    // ProseMirror preserves plugin state by key on reconfigure, even when the
+    // replacement plugin belongs to a different document. Drop selected keys
+    // from an intermediate *state* (not the view), then install the new plugins
+    // in one view update so their state.init runs before their view hooks.
+    const resetKeys = new Set(resetPluginStateFor);
+    const baseState = resetKeys.size
+      ? currentState.reconfigure({
+          plugins: currentState.plugins.filter(
+            (plugin) => !plugin.spec.key || !resetKeys.has(plugin.spec.key),
+          ),
+        })
+      : currentState;
+    const state = baseState.reconfigure({
       plugins: update(currentState.plugins.slice()),
     });
 

packages/core/src/yjs/extensions/ForkYDoc.test.ts

@@ -2,6 +2,7 @@ import { afterEach, describe, expect, it } from "vite-plus/test";
 import { trackPosition } from "../../api/positionMapping.js";
 import * as Y from "yjs";
 import { Awareness } from "y-protocols/awareness";
+import { ySyncPluginKey, yUndoPluginKey } from "y-prosemirror";
 import { BlockNoteEditor } from "../../index.js";
 import { ForkYDocExtension } from "./ForkYDoc.js";
 import { withCollaboration } from "./index.js";
@@ -13,18 +14,19 @@ import { withCollaboration } from "./index.js";
 function createCollabEditor() {
   const doc = new Y.Doc();
   const fragment = doc.getXmlFragment("doc");
+  const awareness = new Awareness(doc);
   const editor = BlockNoteEditor.create(
     withCollaboration({
       collaboration: {
         fragment,
         user: { name: "Test User", color: "#FF0000" },
-        provider: { awareness: new Awareness(doc) },
+        provider: { awareness },
       },
     }),
   );
   const div = document.createElement("div");
   editor.mount(div);
-  return { editor, doc, fragment };
+  return { editor, doc, fragment, awareness };
 }
 
 function getEditorText(editor: BlockNoteEditor) {
@@ -44,6 +46,7 @@ let ctx: ReturnType<typeof createCollabEditor>;
 
 afterEach(() => {
   ctx?.editor.unmount();
+  ctx?.awareness.destroy();
   ctx?.doc.destroy();
 });
 
@@ -90,6 +93,68 @@ describe("ForkYDocExtension", () => {
     expect(getEditorText(ctx.editor)).toContain("Forked edit");
   });
 
+  // https://github.com/TypeCellOS/BlockNote/issues/3135
+  it.each([false, true])(
+    "merge({ keepChanges: %s }) restores remote cursors without splitting the sync binding",
+    (keepChanges) => {
+      ctx = createCollabEditor();
+      ctx.editor.replaceBlocks(ctx.editor.document, [
+        { type: "paragraph", content: "one" },
+        { type: "paragraph", content: "two" },
+      ]);
+
+      const cursor = Y.relativePositionToJSON(
+        Y.createRelativePositionFromTypeIndex(
+          ctx.fragment,
+          ctx.fragment.length,
+        ),
+      );
+      // The provider's awareness map can receive a remote state without a network connection.
+      ctx.awareness.getStates().set(424242, {
+        user: { name: "Remote", color: "#00FF00" },
+        cursor: { anchor: cursor, head: cursor },
+      });
+
+      const forkYDoc = ctx.editor.getExtension(ForkYDocExtension)!;
+      forkYDoc.fork();
+      setEditorText(ctx.editor, "Forked edit");
+      expect(() => forkYDoc.merge({ keepChanges })).not.toThrow();
+      expect(forkYDoc.store.state.isForked).toBe(false);
+      expect(getEditorText(ctx.editor)).toBe(
+        keepChanges ? "Forked edit" : "onetwo",
+      );
+      const sync = ySyncPluginKey.getState(ctx.editor.prosemirrorState);
+      expect(sync.type).toBe(ctx.fragment);
+      expect(sync.doc).toBe(ctx.doc);
+      expect(sync.binding.type).toBe(ctx.fragment);
+      expect(
+        ctx.editor.domElement?.querySelector(".bn-collaboration-cursor__base"),
+      ).not.toBeNull();
+    },
+  );
+
+  it("keeps fork undo separate from the original undo and redo history", () => {
+    ctx = createCollabEditor();
+    setEditorText(ctx.editor, "First");
+    yUndoPluginKey
+      .getState(ctx.editor.prosemirrorState)!
+      .undoManager.stopCapturing();
+    setEditorText(ctx.editor, "Second");
+    expect(ctx.editor.undo()).toBe(true);
+    expect(getEditorText(ctx.editor)).toBe("First");
+
+    const forkYDoc = ctx.editor.getExtension(ForkYDocExtension)!;
+    forkYDoc.fork();
+    setEditorText(ctx.editor, "Forked");
+    expect(ctx.editor.undo()).toBe(true);
+    expect(getEditorText(ctx.editor)).toBe("First");
+    expect(ctx.fragment.toJSON()).toContain("First");
+    forkYDoc.merge({ keepChanges: false });
+
+    expect(ctx.editor.redo()).toBe(true);
+    expect(getEditorText(ctx.editor)).toBe("Second");
+  });
+
   it("fork({ initialUpdate }) uses the provided update instead of the live doc", () => {
     ctx = createCollabEditor();
     setEditorText(ctx.editor, "Current content");

packages/core/src/yjs/extensions/ForkYDoc.ts

@@ -1,6 +1,5 @@
 import { ySyncPluginKey, yUndoPluginKey } from "y-prosemirror";
 import * as Y from "yjs";
-import type { BlockNoteEditor } from "../../editor/BlockNoteEditor.js";
 import {
   createExtension,
   createStore,
@@ -12,32 +11,13 @@ import { YSyncExtension } from "./YSync.js";
 import { YUndoExtension } from "./YUndo.js";
 import { findTypeInOtherYdoc } from "../utils.js";
 
-/**
- * Point the `ySync` plugin state at `fragment`.
- *
- * Swapping the `ySync` plugin reconfigures the ProseMirror state, and
- * ProseMirror carries over the state of plugins that share a key instead of
- * re-initializing them. So the new plugin's `binding` (which is set from its
- * view, via a transaction) ends up on the new fragment, while `type` and `doc`
- * still point at the fragment the editor was bound to before. Anything reading
- * those (e.g. `RelativePositionMappingExtension`) would then mix up the two
- * Y.Docs, so we set them explicitly here.
- */
-function bindYSyncPluginStateTo(
-  editor: BlockNoteEditor<any, any, any>,
-  fragment: Y.XmlFragment,
-) {
-  editor.transact((tr) =>
-    tr.setMeta(ySyncPluginKey, { type: fragment, doc: fragment.doc }),
-  );
-}
-
 export const ForkYDocExtension = createExtension(
   ({ editor, options }: ExtensionOptions<CollaborationOptions>) => {
     let forkedState:
       | {
           originalFragment: Y.XmlFragment;
           undoStack: Y.UndoManager["undoStack"];
+          redoStack: Y.UndoManager["redoStack"];
           forkedFragment: Y.XmlFragment;
         }
       | undefined = undefined;
@@ -81,9 +61,12 @@ export const ForkYDocExtension = createExtension(
         // Find the forked fragment in the new Yjs document
         const forkedFragment = findTypeInOtherYdoc(originalFragment, doc);
 
+        const originalUndoManager = yUndoPluginKey.getState(
+          editor.prosemirrorState,
+        )!.undoManager;
         forkedState = {
-          undoStack: yUndoPluginKey.getState(editor.prosemirrorState)!
-            .undoManager.undoStack,
+          undoStack: originalUndoManager.undoStack,
+          redoStack: originalUndoManager.redoStack,
           originalFragment,
           forkedFragment,
         };
@@ -103,10 +86,9 @@ export const ForkYDocExtension = createExtension(
             // No need to register the cursor plugin again, it's a local fork
             YUndoExtension(),
           ],
+          { resetPluginStateFor: [ySyncPluginKey, yUndoPluginKey] },
         );
 
-        bindYSyncPluginStateTo(editor, forkedFragment);
-
         // Tell the store that the editor is now forked
         store.setState({ isForked: true });
       },
@@ -121,7 +103,8 @@ export const ForkYDocExtension = createExtension(
           return;
         }
 
-        const { originalFragment, forkedFragment, undoStack } = forkedState;
+        const { originalFragment, forkedFragment, undoStack, redoStack } =
+          forkedState;
 
         // Atomically swap the forked plugins back to the original ones
         editor.replaceExtension(
@@ -131,22 +114,23 @@ export const ForkYDocExtension = createExtension(
             YCursorExtension(options),
             YUndoExtension(),
           ],
+          { resetPluginStateFor: [ySyncPluginKey, yUndoPluginKey] },
         );
 
-        bindYSyncPluginStateTo(editor, originalFragment);
-
-        // Reset the undo stack to the original undo stack
-        yUndoPluginKey.getState(
+        // Restore history saved before forking onto the new original-doc manager.
+        const undoManager = yUndoPluginKey.getState(
           editor.prosemirrorState,
-        )!.undoManager.undoStack = undoStack;
+        )!.undoManager;
+        undoManager.undoStack = undoStack;
+        undoManager.redoStack = redoStack;
 
         if (keepChanges) {
           // Apply any changes that have been made to the fork, onto the original doc
           const update = Y.encodeStateAsUpdate(
             forkedFragment.doc!,
             Y.encodeStateVector(originalFragment.doc!),
           );
-          // Applying this change will add to the undo stack, allowing it to be undone normally
+          // Keep the existing editor origin for the merged update.
           Y.applyUpdate(originalFragment.doc!, update, editor);
         }
         // Reset the forked state