feat: AI menu auto scrolling

docs/app/styles.css

@@ -49,14 +49,6 @@ body {
   box-shadow: unset !important;
 }
 
-.demo {
-  overflow: none;
-}
-
-.demo .bn-container {
-  position: relative;
-}
-
 .demo .bn-container:not(.bn-comment-editor),
 .demo .bn-editor {
   height: 100%;
@@ -69,8 +61,7 @@ body {
 
 .demo .bn-editor {
   overflow: auto;
-  padding-top: 1rem;
-  padding-bottom: 250px;
+  padding-block: 1rem;
 }
 
 .demo .bn-editor a {

examples/09-ai/01-minimal/src/App.tsx

@@ -77,6 +77,7 @@ export default function App() {
         // We're disabling some default UI elements
         formattingToolbar={false}
         slashMenu={false}
+        style={{ paddingBottom: "300px" }}
       >
         {/* Add the AI Command menu to the editor */}
         <AIMenuController />

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

@@ -122,6 +122,7 @@ export default function App() {
         editor={editor}
         formattingToolbar={false}
         slashMenu={false}
+        style={{ paddingBottom: "300px" }}
       >
         {/* Add the AI Command menu to the editor */}
         <AIMenuController />

examples/09-ai/03-custom-ai-menu-items/src/App.tsx

@@ -78,6 +78,7 @@ export default function App() {
         editor={editor}
         formattingToolbar={false}
         slashMenu={false}
+        style={{ paddingBottom: "300px" }}
       >
         {/* Creates a new AIMenu with the default items, 
         as well as our custom ones. */}

packages/xl-ai/src/AIExtension.ts

@@ -65,6 +65,9 @@ export class AIExtension extends BlockNoteExtension {
       }
     | undefined;
 
+  private scrollInProgress = false;
+  private autoScroll = false;
+
   public static key(): string {
     return "ai";
   }
@@ -134,6 +137,51 @@ export class AIExtension extends BlockNoteExtension {
         options.agentCursor || { name: "AI", color: "#8bc6ff" },
       ),
     );
+
+    // Scrolls to the block being edited by the AI while auto scrolling is
+    // enabled.
+    this.editor.onCreate(() => {
+      this.editor.onChange(() => {
+        if (!this.autoScroll) {
+          return;
+        }
+
+        const aiMenuState = this._store.getState().aiMenuState;
+        const aiMenuNonErrorState =
+          aiMenuState === "closed" ? undefined : aiMenuState;
+        if (aiMenuNonErrorState?.status === "ai-writing") {
+          const blockElement = this.editor.domElement?.querySelector(
+            `[data-node-type="blockContainer"][data-id="${aiMenuNonErrorState.blockId}"]`,
+          );
+          blockElement?.scrollIntoView({ block: "center" });
+        }
+      });
+    });
+
+    // Listens for `scroll` and `scrollend` events to see if a new scroll was
+    // started before an existing one ended. This is the most reliable way we
+    // have of checking if a scroll event was caused by the user and not by
+    // `scrollIntoView`, as the events are otherwise indistinguishable. If a
+    // scroll was started before an existing one finished (meaning the user has
+    // scrolled), auto scrolling is disabled.
+    document.addEventListener(
+      "scroll",
+      () => {
+        if (this.scrollInProgress) {
+          this.autoScroll = false;
+        }
+
+        this.scrollInProgress = true;
+      },
+      true,
+    );
+    document.addEventListener(
+      "scrollend",
+      () => {
+        this.scrollInProgress = false;
+      },
+      true,
+    );
   }
 
   /**
@@ -148,6 +196,12 @@ export class AIExtension extends BlockNoteExtension {
         status: "user-input",
       },
     });
+
+    // Scrolls to the block when the menu opens.
+    const blockElement = this.editor.domElement?.querySelector(
+      `[data-node-type="blockContainer"][data-id="${blockID}"]`,
+    );
+    blockElement?.scrollIntoView({ block: "center" });
   }
 
   /**
@@ -387,6 +441,7 @@ export class AIExtension extends BlockNoteExtension {
         sender,
         chatRequestOptions: opts.chatRequestOptions,
         onStart: () => {
+          this.autoScroll = true;
           this.setAIResponseStatus("ai-writing");
         },
       });

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

@@ -1,5 +1,5 @@
 import { useBlockNoteEditor } from "@blocknote/react";
-import { FC, useEffect, useState } from "react";
+import { FC } from "react";
 import { useStore } from "zustand";
 
 import { getAIExtension } from "../../AIExtension.js";
@@ -16,82 +16,6 @@ export const AIMenuController = (props: { aiMenu?: FC<AIMenuProps> }) => {
 
   const Component = props.aiMenu || AIMenu;
 
-  const [aiWriting, setAiWriting] = useState(false);
-  const [autoScroll, setAutoScroll] = useState(false);
-  const [scrollInProgress, setScrollInProgress] = useState(false);
-
-  // Converts the `aiMenuState` status to a boolean which shows if the AI is
-  // writing or not. This allows for proper reactivity in other `useEffect`
-  // hooks, while using the base `aiMenuState` object would constantly
-  // retrigger them.
-  useEffect(() => {
-    if (
-      typeof aiMenuState === "object" &&
-      "status" in aiMenuState &&
-      aiMenuState.status === "ai-writing"
-    ) {
-      setAiWriting(true);
-    } else {
-      setAiWriting(false);
-    }
-  }, [aiMenuState]);
-
-  // Enables auto scrolling when the AI starts writing and disables it when it
-  // stops writing.
-  useEffect(() => {
-    if (aiWriting) {
-      setAutoScroll(true);
-    } else {
-      setAutoScroll(false);
-    }
-  }, [aiWriting]);
-
-  // Scrolls to the block being edited by the AI while auto scrolling is
-  // enabled.
-  useEffect(() => {
-    const scrollToBottom = () => {
-      if (!autoScroll) {
-        return;
-      }
-
-      const blockElement = editor.domElement?.querySelector(
-        `[data-node-type="blockContainer"][data-id="${blockId}"]`,
-      );
-      blockElement?.scrollIntoView({ block: "center" });
-    };
-
-    const destroy = editor.onChange(scrollToBottom);
-
-    return () => destroy();
-  }, [autoScroll, blockId, editor]);
-
-  // Listens for `scroll` and `scrollend` events to see if a new scroll was
-  // started before an existing one ended. This is the most reliable way we
-  // have of checking if a scroll event was caused by the user and not by
-  // `scrollIntoView`, as the events are otherwise indistinguishable. If a
-  // scroll was started before an existing one finished (meaning the user has
-  // scrolled), auto scrolling is disabled.
-  useEffect(() => {
-    const scrollHandler = () => {
-      if (scrollInProgress) {
-        setAutoScroll(false);
-      }
-
-      setScrollInProgress(true);
-    };
-    const scrollEndHandler = () => setScrollInProgress(false);
-
-    // Need to set capture to `true` so the events get handled regardless of
-    // which element gets scrolled.
-    document.addEventListener("scroll", scrollHandler, true);
-    document.addEventListener("scrollend", scrollEndHandler, true);
-
-    return () => {
-      document.removeEventListener("scroll", scrollHandler, true);
-      document.removeEventListener("scrollend", scrollEndHandler, true);
-    };
-  }, [scrollInProgress]);
-
   return (
     <BlockPositioner
       canDismissViaOutsidePress={