fix(ai): operations on collaborative documents

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

@@ -1,4 +1,5 @@
 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 { BlockNoteEditor } from "../../index.js";
@@ -213,4 +214,40 @@ describe("ForkYDocExtension", () => {
     forkYDoc.merge({ keepChanges: true });
     expect(getEditorText(ctx.editor)).toContain("Forked modification");
   });
+
+  // https://github.com/TypeCellOS/BlockNote/issues/2946
+  it("can track positions while forked", () => {
+    ctx = createCollabEditor();
+    setEditorText(ctx.editor, "Hello World");
+
+    const forkYDoc = ctx.editor.getExtension(ForkYDocExtension)!;
+    forkYDoc.fork();
+
+    // Store position at "Hello| World"
+    const getCursorPos = trackPosition(ctx.editor, 8);
+    expect(getCursorPos()).toBe(8);
+
+    // Insert text at the beginning of "|Hello World"
+    ctx.editor._tiptapEditor.commands.insertContentAt(3, "Test ");
+    expect(getCursorPos()).toBe(13);
+  });
+
+  // https://github.com/TypeCellOS/BlockNote/issues/2946
+  it("can track positions across fork and merge", () => {
+    ctx = createCollabEditor();
+    setEditorText(ctx.editor, "Hello World");
+
+    // Store position at "Hello| World"
+    const getCursorPos = trackPosition(ctx.editor, 8);
+
+    const forkYDoc = ctx.editor.getExtension(ForkYDocExtension)!;
+    forkYDoc.fork();
+    expect(getCursorPos()).toBe(8);
+
+    ctx.editor._tiptapEditor.commands.insertContentAt(3, "Test ");
+    expect(getCursorPos()).toBe(13);
+
+    forkYDoc.merge({ keepChanges: true });
+    expect(getCursorPos()).toBe(13);
+  });
 });

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

@@ -1,5 +1,6 @@
-import { yUndoPluginKey } from "y-prosemirror";
+import { ySyncPluginKey, yUndoPluginKey } from "y-prosemirror";
 import * as Y from "yjs";
+import type { BlockNoteEditor } from "../../editor/BlockNoteEditor.js";
 import {
   createExtension,
   createStore,
@@ -11,6 +12,26 @@ 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:
@@ -84,6 +105,8 @@ export const ForkYDocExtension = createExtension(
           ],
         );
 
+        bindYSyncPluginStateTo(editor, forkedFragment);
+
         // Tell the store that the editor is now forked
         store.setState({ isForked: true });
       },
@@ -110,6 +133,8 @@ export const ForkYDocExtension = createExtension(
           ],
         );
 
+        bindYSyncPluginStateTo(editor, originalFragment);
+
         // Reset the undo stack to the original undo stack
         yUndoPluginKey.getState(
           editor.prosemirrorState,

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

@@ -51,9 +51,15 @@ export const RelativePositionMappingExtension = createExtension(
           const curYSyncPluginState = ySyncPluginKey.getState(
             editor.prosemirrorState,
           ) as typeof ySyncPluginState;
+          // Resolve against the doc that owns the currently bound type, and not
+          // against `curYSyncPluginState.doc`. Those can point at different
+          // Y.Docs (e.g. right after forking the doc, see `ForkYDocExtension`),
+          // in which case the resolved type wouldn't be part of the bound
+          // fragment and the position would be reported as "not found".
+          const boundType = curYSyncPluginState.binding.type;
           const pos = relativePositionToAbsolutePosition(
-            curYSyncPluginState.doc,
-            curYSyncPluginState.binding.type,
+            boundType.doc,
+            boundType,
             relativePosition,
             curYSyncPluginState.binding.mapping,
           );

packages/xl-ai/package.json

@@ -114,7 +114,8 @@
     "typescript": "^5.9.3",
     "undici": "^6.22.0",
     "vite-plugin-externalize-deps": "^0.10.0",
-    "vite-plus": "catalog:"
+    "vite-plus": "catalog:",
+    "yjs": "^13.6.27"
   },
   "peerDependencies": {
     "react": "^18.0 || ^19.0 || >= 19.0.0-rc",

packages/xl-ai/src/api/formats/html-blocks/collabUpdate.test.ts

@@ -0,0 +1,137 @@
+/**
+ * Regression test for https://github.com/TypeCellOS/BlockNote/issues/2946
+ *
+ * Runs the `update` stream tool against a Yjs-collaborative editor, fully
+ * offline (no LLM call): the tool call is fed straight into the executor.
+ *
+ * `AIExtension.invokeAI` forks the Y.Doc before the request starts, so the
+ * fork is part of the setup here.
+ */
+import {
+  BlockNoteEditor,
+  expandPMRangeToWords,
+  getBlockInfo,
+  getNodeById,
+} from "@blocknote/core";
+import type { ForkYDocExtension } from "@blocknote/core/yjs";
+import { withCollaboration } from "@blocknote/core/yjs";
+import { TextSelection } from "prosemirror-state";
+import { describe, expect, it } from "vite-plus/test";
+import * as Y from "yjs";
+
+import { AIExtension } from "../../../AIExtension.js";
+import { StreamToolExecutor } from "../../../streamTool/StreamToolExecutor.js";
+import { StreamTool } from "../../../streamTool/streamTool.js";
+import { tools } from "./tools/index.js";
+
+function createLocalEditor(text: string) {
+  const editor = BlockNoteEditor.create({
+    initialContent: [{ type: "paragraph", content: text }],
+    trailingBlock: false,
+    extensions: [AIExtension()],
+  });
+  editor.mount(document.createElement("div"));
+
+  return { editor, fork: () => undefined };
+}
+
+function createCollabEditor(text: string) {
+  const ydoc = new Y.Doc();
+  const editor = BlockNoteEditor.create(
+    withCollaboration({
+      collaboration: {
+        fragment: ydoc.getXmlFragment("doc"),
+        user: { color: "#ff0000", name: "Local User" },
+        provider: undefined,
+      },
+      trailingBlock: false,
+      extensions: [AIExtension()],
+    }),
+  );
+  editor.mount(document.createElement("div"));
+
+  editor.replaceBlocks(editor.document, [{ type: "paragraph", content: text }]);
+
+  return {
+    editor,
+    fork: () =>
+      editor.getExtension<typeof ForkYDocExtension>("yForkDoc")?.fork(),
+  };
+}
+
+/**
+ * Selects the full content of the first block, mirroring what the AI menu does
+ * (`buildAIRequest` -> `expandPMRangeToWords`)
+ */
+function selectWholeFirstBlock(editor: BlockNoteEditor<any, any, any>) {
+  const id = editor.document[0].id;
+  const info = getBlockInfo(getNodeById(id, editor.prosemirrorState.doc)!);
+  if (!info.isBlockContainer) {
+    throw new Error("not a block container");
+  }
+  const from = info.blockContent.beforePos + 1;
+  const to = info.blockContent.afterPos - 1;
+
+  editor.transact((tr) => {
+    tr.setSelection(TextSelection.create(tr.doc, from, to));
+  });
+
+  return expandPMRangeToWords(editor.prosemirrorState.doc, {
+    $from: editor.prosemirrorState.doc.resolve(from),
+    $to: editor.prosemirrorState.doc.resolve(to),
+  });
+}
+
+async function runUpdate(
+  editor: BlockNoteEditor<any, any, any>,
+  id: string,
+  html: string,
+  selection?: { from: number; to: number },
+) {
+  const streamTools = [
+    tools.update(editor, {
+      idsSuffixed: false,
+      withDelays: false,
+      updateSelection: selection,
+    }),
+  ] as StreamTool<any>[];
+
+  await new StreamToolExecutor(streamTools).execute(
+    (async function* () {
+      yield {
+        operation: { type: "update" as const, id, block: html },
+        isUpdateToPreviousOperation: false,
+        isPossiblyPartial: false,
+        metadata: undefined,
+      };
+    })(),
+  );
+}
+
+describe.each([
+  ["local", createLocalEditor],
+  ["collaborative", createCollabEditor],
+])("update tool (%s)", (_name, createEditor) => {
+  it("updates a selected paragraph", async () => {
+    const { editor, fork } = createEditor("Bonjour le monde");
+    fork();
+    const id = editor.document[0].id;
+    const selection = selectWholeFirstBlock(editor);
+
+    await runUpdate(editor, id, "<p>Bonjour à tous</p>", selection);
+
+    editor.getExtension(AIExtension)?.acceptChanges();
+    expect((editor.document[0] as any).content[0].text).toBe("Bonjour à tous");
+  });
+
+  it("updates a paragraph without a selection", async () => {
+    const { editor, fork } = createEditor("Bonjour le monde");
+    fork();
+    const id = editor.document[0].id;
+
+    await runUpdate(editor, id, "<p>Bonjour à tous</p>");
+
+    editor.getExtension(AIExtension)?.acceptChanges();
+    expect((editor.document[0] as any).content[0].text).toBe("Bonjour à tous");
+  });
+});