fix(ai): loosen serialization of blocks in columns (#2716) (#2718)

Author: nperez0111

packages/core/src/api/blockManipulation/commands/insertBlocks/insertBlocks.ts

@@ -26,9 +26,11 @@ export function insertBlocks<
   const id =
     typeof referenceBlock === "string" ? referenceBlock : referenceBlock.id;
   const pmSchema = getPmSchema(tr);
-  const nodesToInsert = blocksToInsert.map((block) =>
-    blockToNode(block, pmSchema),
-  );
+  const nodesToInsert = blocksToInsert.map((block) => {
+    const node = blockToNode(block, pmSchema);
+    node.check(); // `blockToNode` is lenient; validate before mutating the doc
+    return node;
+  });
 
   const posInfo = getNodeById(id, tr.doc);
   if (!posInfo) {

packages/core/src/api/blockManipulation/commands/replaceBlocks/replaceBlocks.ts

@@ -27,9 +27,11 @@ export function removeAndInsertBlocks<
   const pmSchema = getPmSchema(tr);
   // Converts the `PartialBlock`s to ProseMirror nodes to insert them into the
   // document.
-  const nodesToInsert: Node[] = blocksToInsert.map((block) =>
-    blockToNode(block, pmSchema),
-  );
+  const nodesToInsert: Node[] = blocksToInsert.map((block) => {
+    const node = blockToNode(block, pmSchema);
+    node.check(); // `blockToNode` is lenient; validate before mutating the doc
+    return node;
+  });
 
   const idsOfBlocksToRemove = new Set<string>(
     blocksToRemove.map((block) =>

packages/core/src/api/blockManipulation/commands/updateBlock/updateBlock.ts

@@ -128,16 +128,18 @@ export function updateBlockTr<
     // for this, we do a nodeToBlock on the existing block to get the children.
     // it would be cleaner to use a ReplaceAroundStep, but this is a bit simpler and it's quite an edge case
     const existingBlock = nodeToBlock(blockInfo.bnBlock.node, pmSchema);
+    const replacementNode = blockToNode(
+      {
+        children: existingBlock.children, // if no children are passed in, use existing children
+        ...block,
+      },
+      pmSchema,
+    );
+    replacementNode.check(); // `blockToNode` is lenient; validate before mutating the doc
     tr.replaceWith(
       blockInfo.bnBlock.beforePos,
       blockInfo.bnBlock.afterPos,
-      blockToNode(
-        {
-          children: existingBlock.children, // if no children are passed in, use existing children
-          ...block,
-        },
-        pmSchema,
-      ),
+      replacementNode,
     );
 
     return;
@@ -278,7 +280,9 @@ function updateChildren<
   const pmSchema = getPmSchema(tr);
   if (block.children !== undefined && block.children.length > 0) {
     const childNodes = block.children.map((child) => {
-      return blockToNode(child, pmSchema);
+      const node = blockToNode(child, pmSchema);
+      node.check(); // `blockToNode` is lenient; validate before mutating the doc
+      return node;
     });
 
     // Checks if a blockGroup node already exists.

packages/core/src/api/nodeConversions/blockToNode.ts

@@ -366,8 +366,9 @@ export function blockToNode(
       groupNode ? [contentNode, groupNode] : contentNode,
     );
   } else if (schema.nodes[block.type].isInGroup("bnBlock")) {
-    // this is a bnBlock node like Column or ColumnList that directly translates to a prosemirror node
-    return schema.nodes[block.type].createChecked(
+    // `create` (not `createChecked`) so partial container blocks pass through;
+    // callers that mutate the doc validate via `node.check()` before inserting.
+    return schema.nodes[block.type].create(
       {
         id: id,
         ...block.props,

packages/xl-ai/package.json

@@ -99,6 +99,7 @@
     "@ai-sdk/mistral": "^3.0.2",
     "@ai-sdk/openai": "^3.0.2",
     "@ai-sdk/openai-compatible": "^2.0.2",
+    "@blocknote/xl-multi-column": "0.50.0",
     "@mswjs/interceptors": "^0.37.6",
     "@types/diff": "^6.0.0",
     "@types/json-diff": "^1.0.3",

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

@@ -0,0 +1,106 @@
+// @vitest-environment jsdom
+
+import { BlockNoteEditor, BlockNoteSchema } from "@blocknote/core";
+import { withMultiColumn } from "@blocknote/xl-multi-column";
+import { afterEach, beforeEach, describe, expect, it } from "vitest";
+
+import { htmlBlockLLMFormat } from "./htmlBlocks.js";
+
+const schema = withMultiColumn(BlockNoteSchema.create());
+
+// Regression test for https://github.com/TypeCellOS/BlockNote/issues/2716
+describe("htmlBlockLLMFormat.defaultDocumentStateBuilder with multi-column", () => {
+  let editor: BlockNoteEditor<any, any, any>;
+  const div = document.createElement("div");
+
+  beforeEach(() => {
+    editor = BlockNoteEditor.create({
+      schema,
+      initialContent: [
+        {
+          id: "column-list-0",
+          type: "columnList",
+          children: [
+            {
+              id: "column-0",
+              type: "column",
+              children: [
+                {
+                  id: "left-paragraph",
+                  type: "paragraph",
+                  content: "left column",
+                },
+              ],
+            },
+            {
+              id: "column-1",
+              type: "column",
+              children: [
+                {
+                  id: "right-paragraph",
+                  type: "paragraph",
+                  content: "right column",
+                },
+              ],
+            },
+          ],
+        },
+      ],
+    });
+    editor.mount(div);
+  });
+
+  afterEach(() => {
+    editor._tiptapEditor.destroy();
+    editor = undefined as any;
+  });
+
+  it("builds the document state for a doc containing a columnList without throwing", async () => {
+    const result = await htmlBlockLLMFormat.defaultDocumentStateBuilder({
+      editor,
+      streamTools: [],
+      onStart: () => {
+        // no-op
+      },
+    });
+
+    if (result.selection !== false) {
+      throw new Error("expected selection-less document state");
+    }
+
+    const idEntries = result.blocks.filter(
+      (b): b is { id: string; block: string } => "id" in b,
+    );
+
+    const columnListEntry = idEntries.find((b) =>
+      b.id.includes("column-list-0"),
+    );
+    expect(columnListEntry).toBeDefined();
+    expect(columnListEntry!.block).toContain('data-node-type="columnList"');
+
+    expect(idEntries.find((b) => b.id.includes("column-0"))).toBeDefined();
+    expect(idEntries.find((b) => b.id.includes("column-1"))).toBeDefined();
+    expect(
+      idEntries.find((b) => b.id.includes("left-paragraph")),
+    ).toBeDefined();
+    expect(
+      idEntries.find((b) => b.id.includes("right-paragraph")),
+    ).toBeDefined();
+  });
+
+  it("builds the document state when invoked with selectedBlocks containing a columnList", async () => {
+    const result = await htmlBlockLLMFormat.defaultDocumentStateBuilder({
+      editor,
+      selectedBlocks: editor.document,
+      streamTools: [],
+      onStart: () => {
+        // no-op
+      },
+    });
+
+    if (result.selection !== true) {
+      throw new Error("expected selection-based document state");
+    }
+    expect(result.selectedBlocks.length).toBeGreaterThan(0);
+  });
+});

packages/xl-ai/vite.config.ts

@@ -26,6 +26,10 @@ export default defineConfig((conf) => ({
             "@blocknote/core": path.resolve(__dirname, "../core/src/"),
             "@blocknote/mantine": path.resolve(__dirname, "../mantine/src/"),
             "@blocknote/react": path.resolve(__dirname, "../react/src/"),
+            "@blocknote/xl-multi-column": path.resolve(
+              __dirname,
+              "../xl-multi-column/src/",
+            ),
             "@shared": path.resolve(__dirname, "../../shared/"),
           } as Record<string, string>),
   },