fix(ai): operations on blocks containing comments

packages/core/src/yjs/comments/RESTYjsThreadStore.test.ts

@@ -0,0 +1,85 @@
+/**
+ * @vitest-environment jsdom
+ */
+import {
+  relativePositionToAbsolutePosition,
+  ySyncPluginKey,
+} from "y-prosemirror";
+import { afterEach, describe, expect, it, vi } from "vite-plus/test";
+import * as Y from "yjs";
+import { BlockNoteEditor } from "../../editor/BlockNoteEditor.js";
+import { DefaultThreadStoreAuth } from "../../comments/threadstore/DefaultThreadStoreAuth.js";
+import { withCollaboration } from "../extensions/index.js";
+import { RESTYjsThreadStore } from "./RESTYjsThreadStore.js";
+
+function createCollabEditor() {
+  const doc = new Y.Doc();
+  const fragment = doc.getXmlFragment("doc");
+  const editor = BlockNoteEditor.create(
+    withCollaboration({
+      collaboration: {
+        fragment,
+        user: { name: "Test User", color: "#FF0000" },
+        provider: undefined,
+      },
+      trailingBlock: false,
+    }),
+  );
+  editor.mount(document.createElement("div"));
+  editor.replaceBlocks(editor.document, [
+    { type: "paragraph", content: "Hello World" },
+  ]);
+  return { editor, doc, fragment };
+}
+
+afterEach(() => {
+  vi.restoreAllMocks();
+});
+
+describe("RESTYjsThreadStore", () => {
+  it("sends resolvable yjs positions along with the thread", async () => {
+    const { editor, doc, fragment } = createCollabEditor();
+
+    const requests: any[] = [];
+    vi.spyOn(globalThis, "fetch").mockImplementation((async (
+      _url: any,
+      init: any,
+    ) => {
+      requests.push(JSON.parse(init.body));
+      return new Response("{}", {
+        status: 200,
+        headers: { "Content-Type": "application/json" },
+      });
+    }) as any);
+
+    const store = new RESTYjsThreadStore(
+      "https://example.com/threads",
+      {},
+      doc.getMap("threads"),
+      new DefaultThreadStoreAuth("user-1", "editor"),
+    );
+
+    await store.addThreadToDocument({
+      threadId: "thread-1",
+      selection: { anchor: 3, head: 8 },
+      editor,
+    });
+
+    expect(requests).toHaveLength(1);
+    const { yjs } = requests[0].selection;
+    expect(yjs).toBeDefined();
+
+    // the relative positions must resolve back to the positions we passed in
+    const state = ySyncPluginKey.getState(editor.prosemirrorState) as any;
+    const resolve = (relPos: any) =>
+      relativePositionToAbsolutePosition(
+        fragment.doc!,
+        state.binding.type,
+        Y.createRelativePositionFromJSON(relPos),
+        state.binding.mapping,
+      );
+
+    expect(resolve(yjs.anchor)).toBe(3);
+    expect(resolve(yjs.head)).toBe(8);
+  });
+});

packages/core/src/yjs/comments/RESTYjsThreadStore.ts

@@ -59,7 +59,13 @@ export class RESTYjsThreadStore extends YjsThreadStoreBase {
   }) => {
     const { threadId, selection } = options;
 
-    const binding = ySyncPluginKey.getState(options.editor.prosemirrorState);
+    // Note: the positions have to be resolved against the *binding's* type and
+    // mapping. The plugin state has a `type` of its own, but no `mapping`, and
+    // its `type` can go stale (e.g. while the doc is forked, see
+    // `ForkYDocExtension`).
+    const binding = ySyncPluginKey.getState(
+      options.editor.prosemirrorState,
+    )?.binding;
     const yjsSelection = binding
       ? {
           head: absolutePositionToRelativePosition(

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

@@ -0,0 +1,122 @@
+/**
+ * Regression test for https://github.com/TypeCellOS/BlockNote/issues/2947
+ *
+ * Comment marks are `blocknoteIgnore` marks: they're not part of the BlockNote
+ * document model and don't survive an HTML round-trip. The HTML rebase tool
+ * used to treat that as a fatal "html diff", which made every AI operation on a
+ * block containing a comment fail.
+ *
+ * Runs fully offline (no LLM call): tool calls are fed straight into the
+ * executor.
+ */
+import { BlockNoteEditor, createExtension } from "@blocknote/core";
+import { CommentMark } from "@blocknote/core/comments";
+import { describe, expect, it } from "vite-plus/test";
+
+import { AIExtension } from "../../../AIExtension.js";
+import { StreamToolExecutor } from "../../../streamTool/StreamToolExecutor.js";
+import { StreamTool } from "../../../streamTool/streamTool.js";
+import { tools } from "./tools/index.js";
+import { createHTMLRebaseTool } from "./tools/rebaseTool.js";
+
+/**
+ * Registers just the comment mark. The full `CommentsExtension` needs a thread
+ * store and user resolver, neither of which affects the document model.
+ */
+const CommentMarkExtension = createExtension(() => ({
+  key: "commentMarkOnly",
+  tiptapExtensions: [CommentMark],
+}));
+
+function createEditorWithComment() {
+  const editor = BlockNoteEditor.create({
+    initialContent: [
+      { id: "ref1", type: "paragraph", content: "Hello, world!" },
+      { id: "ref2", type: "paragraph", content: "How are you?" },
+    ],
+    trailingBlock: false,
+    extensions: [AIExtension(), CommentMarkExtension()],
+  });
+  editor.mount(document.createElement("div"));
+
+  // comment on "Hello" in the first block
+  editor.transact((tr) =>
+    tr.addMark(
+      3,
+      8,
+      editor.pmSchema.marks.comment.create({
+        threadId: "thread-1",
+        orphan: false,
+      }),
+    ),
+  );
+
+  return editor;
+}
+
+function commentedRanges(editor: BlockNoteEditor<any, any, any>) {
+  const ranges: { text: string; threadId: string }[] = [];
+  editor.prosemirrorState.doc.descendants((node) => {
+    const mark = node.marks.find((m) => m.type.name === "comment");
+    if (mark) {
+      ranges.push({ text: node.text!, threadId: mark.attrs.threadId });
+    }
+  });
+  return ranges;
+}
+
+async function runUpdate(
+  editor: BlockNoteEditor<any, any, any>,
+  id: string,
+  html: string,
+) {
+  const streamTools = [
+    tools.update(editor, { idsSuffixed: false, withDelays: false }),
+  ] 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("blocks containing comments", () => {
+  it("can build a rebase tool for a commented block", () => {
+    const editor = createEditorWithComment();
+
+    expect(() => createHTMLRebaseTool("ref1", editor)).not.toThrow();
+  });
+
+  it("updates a block that contains a comment", async () => {
+    const editor = createEditorWithComment();
+
+    await runUpdate(editor, "ref1", "<p>Hello, universe!</p>");
+
+    editor.getExtension(AIExtension)?.acceptChanges();
+    expect((editor.document[0] as any).content[0].text).toBe(
+      "Hello, universe!",
+    );
+    // the untouched part of the comment is still anchored
+    expect(commentedRanges(editor)).toEqual([
+      { text: "Hello", threadId: "thread-1" },
+    ]);
+  });
+
+  it("updates a sibling block while another block has a comment", async () => {
+    const editor = createEditorWithComment();
+
+    await runUpdate(editor, "ref2", "<p>How do you do?</p>");
+
+    editor.getExtension(AIExtension)?.acceptChanges();
+    expect((editor.document[1] as any).content[0].text).toBe("How do you do?");
+    expect(commentedRanges(editor)).toEqual([
+      { text: "Hello", threadId: "thread-1" },
+    ]);
+  });
+});

packages/xl-ai/src/api/formats/html-blocks/tools/rebaseTool.ts

@@ -1,4 +1,5 @@
 import { BlockNoteEditor, getBlock } from "@blocknote/core";
+import { Mapping } from "prosemirror-transform";
 import { updateToReplaceSteps } from "../../../../prosemirror/changeset.js";
 import {
   getApplySuggestionsTr,
@@ -47,13 +48,24 @@ export function createHTMLRebaseTool(
     tr.doc,
   );
 
-  if (steps.length) {
-    throw new Error("html diff", {
-      cause: {
-        html,
-        htmlBlock,
-      },
-    });
+  // The HTML round-trip isn't always lossless: marks that aren't part of the
+  // BlockNote document model (`blocknoteIgnore`, e.g. comments) don't survive
+  // it. Apply the difference to the projection, so operations are applied to
+  // the document as the HTML format sees it, and are then rebased onto the
+  // actual document (same as `createMDRebaseTool` does for markdown).
+  const stepMapping = new Mapping();
+  for (const step of steps) {
+    const mapped = step.map(stepMapping);
+    if (!mapped) {
+      throw new Error("html diff", {
+        cause: {
+          html,
+          htmlBlock,
+        },
+      });
+    }
+    tr.step(mapped);
+    stepMapping.appendMap(mapped.getMap());
   }
 
   return rebaseTool(editor, tr);