Refactored/tested/fixed AI menu and other misc changes

Author: matthewlipski

examples/09-ai/02-playground/src/App.tsx

@@ -87,6 +87,7 @@ export default function App() {
 
     // Add the model string to the request body
     ai.options.setState({
+      ...ai.options.state,
       chatRequestOptions: {
         body: {
           model,

packages/react/src/components/LinkToolbar/LinkToolbarController.tsx

@@ -159,6 +159,7 @@ export const LinkToolbarController = (props: {
   // an x, y, height, and width of 0. This is why we instead use a virtual
   // element for the reference, and use the last obtained `DOMRect` from when
   // the link element was still mounted to the DOM.
+  // TODO: Logic should be in `GenericPopover`.
   const mountedBoundingClientRect = useRef(new DOMRect());
   if (link?.element && editor.prosemirrorView.root.contains(link.element)) {
     mountedBoundingClientRect.current = link.element.getBoundingClientRect();

packages/react/src/components/Popovers/BlockPopover.tsx

@@ -1,5 +1,6 @@
 import { getNodeById } from "@blocknote/core";
-import { ReactNode, useMemo } from "react";
+import { ClientRectObject, VirtualElement } from "@floating-ui/react";
+import { ReactNode, useMemo, useRef } from "react";
 
 import { useBlockNoteEditor } from "../../hooks/useBlockNoteEditor.js";
 import { FloatingUIOptions } from "./FloatingUIOptions.js";
@@ -11,41 +12,51 @@ export const BlockPopover = (
     children: ReactNode;
   },
 ) => {
-  const { blockId, children, ...rest } = props;
+  const { blockId, children, ...floatingUIOptions } = props;
 
   const editor = useBlockNoteEditor<any, any, any>();
 
-  // Seems like unlike with virtual elements, we don't need to use the last
-  // defined reference element in case it's currently undefined. FloatingUI
-  // appears to already do this internally.
-  const element = useMemo(
+  // Stores the last `boundingClientRect` to use from when the block's DOM
+  // element was still mounted. This is because `DOMRect`s returned from
+  // calling `getBoundingClientRect` on unmounted elements will have x, y,
+  // height, and width of 0. This can cause issues when the popover is
+  // transitioning out.
+  // TODO: Move this logic to the `GenericPopover`.
+  const mountedBoundingClientRect = useRef<ClientRectObject>(new DOMRect());
+  const virtualElement = useMemo(
     () =>
       editor.transact((tr) => {
+        const virtualElement: VirtualElement = {
+          getBoundingClientRect: () => mountedBoundingClientRect.current,
+        };
+
         if (!blockId) {
-          return undefined;
+          return virtualElement;
         }
 
         // TODO use the location API for this
         const nodePosInfo = getNodeById(blockId, tr.doc);
         if (!nodePosInfo) {
-          return undefined;
+          return virtualElement;
         }
 
         const { node } = editor.prosemirrorView.domAtPos(
           nodePosInfo.posBeforeNode + 1,
         );
         if (!(node instanceof Element)) {
-          return undefined;
+          return virtualElement;
         }
 
-        return node;
+        mountedBoundingClientRect.current = node.getBoundingClientRect();
+
+        return virtualElement;
       }),
     [editor, blockId],
   );
 
   return (
-    <GenericPopover reference={element} {...rest}>
-      {children}
+    <GenericPopover reference={virtualElement} {...floatingUIOptions}>
+      {blockId !== undefined && children}
     </GenericPopover>
   );
 };

packages/react/src/components/Popovers/GenericPopover.tsx

@@ -71,10 +71,9 @@ export const GenericPopover = (
         }
       }
     },
-    // All of `props` is included in the deps, as we also need to run this
-    // effect whenever the passed children change, since it's ultimately the
-    // HTML of the children that we're storing.
-    [status, props],
+    // `props.children` is added to the deps, since it's ultimately the HTML of
+    // the children that we're storing.
+    [status, props.reference, props.children],
   );
 
   if (!isMounted) {

packages/react/src/components/Popovers/PositionPopover.tsx

@@ -1,3 +1,4 @@
+import { ClientRectObject } from "@floating-ui/react";
 import { posToDOMRect } from "@tiptap/core";
 import { ReactNode, useMemo, useRef } from "react";
 
@@ -11,15 +12,16 @@ export const PositionPopover = (
     children: ReactNode;
   },
 ) => {
-  const { position, children, ...rest } = props;
+  const { position, children, ...floatingUIOptions } = props;
   const { from, to } = position || {};
 
   const editor = useBlockNoteEditor<any, any, any>();
   const { headless } = editor;
 
   // Stores the last created `boundingClientRect` to use in case `position` is
   // not defined.
-  const boundingClientRect = useRef<DOMRect>(new DOMRect());
+  // TODO: Move this logic to the `GenericPopover`.
+  const boundingClientRect = useRef<ClientRectObject>(new DOMRect());
   const virtualElement = useMemo(
     () => ({
       getBoundingClientRect: () => {
@@ -47,8 +49,8 @@ export const PositionPopover = (
   );
 
   return (
-    <GenericPopover reference={virtualElement} {...rest}>
-      {children}
+    <GenericPopover reference={virtualElement} {...floatingUIOptions}>
+      {position !== undefined && children}
     </GenericPopover>
   );
 };

packages/react/src/components/Popovers/TableCellPopover.tsx

@@ -13,7 +13,7 @@ export const TableCellPopover = (
     children: ReactNode;
   },
 ) => {
-  const { blockId, colIndex, rowIndex, children, ...rest } = props;
+  const { blockId, colIndex, rowIndex, children, ...floatingUIOptions } = props;
 
   const editor = useBlockNoteEditor<any, any, any>();
 
@@ -50,7 +50,7 @@ export const TableCellPopover = (
   }, [editor, blockId, colIndex, rowIndex]);
 
   return (
-    <GenericPopover reference={element} {...rest}>
+    <GenericPopover reference={element} {...floatingUIOptions}>
       {children}
     </GenericPopover>
   );

packages/react/src/components/SideMenu/SideMenuController.tsx

@@ -1,4 +1,4 @@
-import { SideMenuExtension as SideMenuExtension } from "@blocknote/core/extensions";
+import { SideMenuExtension } from "@blocknote/core/extensions";
 import { FC, useMemo } from "react";
 
 import { useExtensionState } from "../../hooks/useExtension.js";
@@ -47,7 +47,7 @@ export const SideMenuController = (props: {
 
   return (
     <BlockPopover blockId={block?.id} {...floatingUIOptions}>
-      <Component />
+      {block?.id && <Component />}
     </BlockPopover>
   );
 };

packages/react/src/index.ts

@@ -47,6 +47,12 @@ export * from "./components/LinkToolbar/LinkToolbar.js";
 export * from "./components/LinkToolbar/LinkToolbarController.js";
 export * from "./components/LinkToolbar/LinkToolbarProps.js";
 
+export * from "./components/Popovers/BlockPopover.js";
+export * from "./components/Popovers/FloatingUIOptions.js";
+export * from "./components/Popovers/GenericPopover.js";
+export * from "./components/Popovers/PositionPopover.js";
+export * from "./components/Popovers/TableCellPopover.js";
+
 export * from "./components/SideMenu/DefaultButtons/AddBlockButton.js";
 export * from "./components/SideMenu/DefaultButtons/DragHandleButton.js";
 export * from "./components/SideMenu/SideMenu.js";

packages/xl-ai/src/AIExtension.ts

@@ -383,41 +383,38 @@ export const AIExtension = createExtension(
             deleteEmptyCursorBlock: opts.deleteEmptyCursorBlock,
             streamToolsProvider: opts.streamToolsProvider,
             onBlockUpdated: (blockId) => {
-              // NOTE: does this setState with an anon object trigger unnecessary re-renders?
-              store.setState({
-                aiMenuState: {
-                  blockId,
-                  status: "ai-writing",
-                },
-              });
-
-              // Scrolls to the block being edited by the AI while auto scrolling is
-              // enabled.
-              if (!autoScroll) {
-                return;
-              }
-
               const aiMenuState = store.state.aiMenuState;
               const aiMenuOpenState =
                 aiMenuState === "closed" ? undefined : aiMenuState;
               if (!aiMenuOpenState || aiMenuOpenState.status !== "ai-writing") {
                 return;
               }
 
+              // TODO: Sometimes, the updated block doesn't actually exist in
+              // the editor. I don't know why this happens, seems like a bug?
               const nodeInfo = getNodeById(
-                aiMenuOpenState.blockId,
+                blockId,
                 editor.prosemirrorState.doc,
               );
               if (!nodeInfo) {
                 return;
               }
 
-              const blockElement = editor.prosemirrorView.domAtPos(
-                nodeInfo.posBeforeNode + 1,
-              );
-              (blockElement.node as HTMLElement).scrollIntoView({
-                block: "center",
+              store.setState({
+                aiMenuState: {
+                  blockId,
+                  status: "ai-writing",
+                },
               });
+
+              if (autoScroll) {
+                const blockElement = editor.prosemirrorView.domAtPos(
+                  nodeInfo.posBeforeNode + 1,
+                );
+                (blockElement.node as HTMLElement).scrollIntoView({
+                  block: "center",
+                });
+              }
             },
           });
 

packages/xl-ai/src/components/AIMenu/AIMenuController.tsx

@@ -1,12 +1,19 @@
-import { useBlockNoteEditor } from "@blocknote/react";
-import { FC } from "react";
+import {
+  BlockPopover,
+  FloatingUIOptions,
+  useBlockNoteEditor,
+} from "@blocknote/react";
+import { FC, useMemo } from "react";
 import { useExtension, useExtensionState } from "@blocknote/react";
+import { offset, size } from "@floating-ui/react";
 
 import { AIExtension } from "../../AIExtension.js";
 import { AIMenu, AIMenuProps } from "./AIMenu.js";
-import { BlockPositioner } from "./BlockPositioner.js";
 
-export const AIMenuController = (props: { aiMenu?: FC<AIMenuProps> }) => {
+export const AIMenuController = (props: {
+  aiMenu?: FC<AIMenuProps>;
+  floatingUIOptions?: FloatingUIOptions;
+}) => {
   const editor = useBlockNoteEditor();
   const ai = useExtension(AIExtension);
 
@@ -17,29 +24,71 @@ export const AIMenuController = (props: { aiMenu?: FC<AIMenuProps> }) => {
 
   const blockId = aiMenuState === "closed" ? undefined : aiMenuState.blockId;
 
+  const floatingUIOptions = useMemo<FloatingUIOptions>(
+    () => ({
+      useFloatingOptions: {
+        open: aiMenuState !== "closed",
+        placement: "bottom",
+        middleware: [
+          offset(10),
+          // flip(),
+          size({
+            apply({ rects, elements }) {
+              Object.assign(elements.floating.style, {
+                width: `${rects.reference.width}px`,
+              });
+            },
+          }),
+        ],
+        onOpenChange: (open) => {
+          if (open || aiMenuState === "closed") {
+            return;
+          }
+
+          if (aiMenuState.status === "user-input") {
+            ai.closeAIMenu();
+          } else if (
+            aiMenuState.status === "user-reviewing" ||
+            aiMenuState.status === "error"
+          ) {
+            ai.rejectChanges();
+          }
+        },
+      },
+      useDismissProps: {
+        enabled:
+          aiMenuState === "closed" || aiMenuState.status === "user-input",
+        // We should just be able to set `referencePress: true` instead of
+        // using this listener, but for some reason it doesn't seem to trigger.
+        outsidePress: (event) => {
+          if (event.target instanceof Element) {
+            const blockElement = event.target.closest(".bn-block");
+            if (
+              blockElement &&
+              blockElement.getAttribute("data-id") === blockId
+            ) {
+              ai.closeAIMenu();
+            }
+          }
+
+          return true;
+        },
+      },
+      elementProps: {
+        style: {
+          zIndex: 100,
+        },
+      },
+      ...props.floatingUIOptions,
+    }),
+    [ai, aiMenuState, blockId, props.floatingUIOptions],
+  );
+
   const Component = props.aiMenu || AIMenu;
 
   return (
-    <BlockPositioner
-      canDismissViaOutsidePress={
-        aiMenuState === "closed" || aiMenuState.status === "user-input"
-      }
-      blockID={blockId}
-      onOpenChange={(open) => {
-        if (open || aiMenuState === "closed") {
-          return;
-        }
-        if (aiMenuState.status === "user-input") {
-          ai.closeAIMenu();
-        } else if (
-          aiMenuState.status === "user-reviewing" ||
-          aiMenuState.status === "error"
-        ) {
-          ai.rejectChanges();
-        }
-      }}
-    >
-      <Component />
-    </BlockPositioner>
+    <BlockPopover blockId={blockId} {...floatingUIOptions}>
+      {aiMenuState !== "closed" && <Component />}
+    </BlockPopover>
   );
 };