diff --git a/docs/content/docs/features/ai/custom-commands.mdx b/docs/content/docs/features/ai/custom-commands.mdx
index de02f19743..75b5c029db 100644
--- a/docs/content/docs/features/ai/custom-commands.mdx
+++ b/docs/content/docs/features/ai/custom-commands.mdx
@@ -26,8 +26,8 @@ First, we define a new AI command. Let's create one that makes selected text mor
 
 ```tsx
 import {
+  AIExtension,
   AIMenuSuggestionItem,
-  getAIExtension,
   aiDocumentFormats,
 } from "@blocknote/xl-ai";
 
@@ -40,7 +40,7 @@ export const makeInformal = (
   aliases: ["informal", "make informal", "casual"],
   icon: <RiEmotionHappyFill size={18} />,
   onItemClick: async () => {
-    await getAIExtension(editor).invokeAI({
+    await editor.getExtension(AIExtension)?.invokeAI({
       // The prompt to send to the LLM:
       userPrompt: "Give the selected text a more informal (casual) tone",
       // Tell the LLM to specifically use the selected content as context (instead of the whole document)
diff --git a/docs/content/docs/features/ai/reference.mdx b/docs/content/docs/features/ai/reference.mdx
index 2bafd6b647..1197efcc6b 100644
--- a/docs/content/docs/features/ai/reference.mdx
+++ b/docs/content/docs/features/ai/reference.mdx
@@ -52,14 +52,6 @@ type AIRequestHelpers = {
 };
 ```
 
-## `getAIExtension`
-
-Use `getAIExtension` to retrieve the AI extension instance registered to the editor:
-
-```typescript
-function getAIExtension(editor: BlockNoteEditor): AIExtension;
-```
-
 ## `AIExtension`
 
 The `AIExtension` class is the main class for the AI extension. It exposes state and methods to interact with BlockNote's AI features.
diff --git a/docs/content/docs/features/extensions.mdx b/docs/content/docs/features/extensions.mdx
index e2fa8e904e..8b1dcd37c0 100644
--- a/docs/content/docs/features/extensions.mdx
+++ b/docs/content/docs/features/extensions.mdx
@@ -14,7 +14,7 @@ BlockNote includes an extensions system which lets you expand the editor's behav
 
 ## Creating an extension
 
-An extension is an instance of the [`BlockNoteExtension`](https://github.com/TypeCellOS/BlockNote/blob/10cdbfb5f77ef82f3617c0fa1191e0bf5b7358c5/packages/core/src/editor/BlockNoteExtension.ts#L13) class. However, it's recommended for most use cases to create extensions using the `createBlockNoteExtension` function, rather than instanciating the class directly:
+An extension is an instance of the [`BlockNoteExtension`](https://github.com/TypeCellOS/BlockNote/blob/10cdbfb5f77ef82f3617c0fa1191e0bf5b7358c5/packages/core/src/editor/BlockNoteExtension.ts#L13) class. However, it's recommended for most use cases to create extensions using the `createExtension` function, rather than instanciating the class directly:
 
 ```typescript
 type BlockNoteExtensionOptions = {
@@ -43,10 +43,10 @@ const customBlockExtensionOptions: BlockNoteExtensionOptions = {
   tiptapExtensions: ...,
 }
 
-const CustomExtension = createBlockNoteExtension(customBlockExtensionOptions);
+const CustomExtension = createExtension(customBlockExtensionOptions);
 ```
 
-Let's go over the options that can be passed into `createBlockNoteExtension`:
+Let's go over the options that can be passed into `createExtension`:
 
 `key:` The name of the extension.
 
@@ -74,7 +74,7 @@ The `extensions` [editor option](/docs/reference/editor/overview#options) takes
 const editor = useCreateBlockNote({
   extensions: [
     // Add extensions here:
-    createBlockNoteExtension({ ... })
+    createExtension({ ... })
   ],
 });
 ```
@@ -95,7 +95,7 @@ const createCustomBlock = createReactBlockSpec(
   }
   [
     // Add extensions here:
-    createBlockNoteExtension({ ... })
+    createExtension({ ... })
   ],
 });
 ```
diff --git a/examples/01-basic/03-multi-column/src/App.tsx b/examples/01-basic/03-multi-column/src/App.tsx
index c126a9d2bf..3a0d10d8d9 100644
--- a/examples/01-basic/03-multi-column/src/App.tsx
+++ b/examples/01-basic/03-multi-column/src/App.tsx
@@ -1,8 +1,8 @@
 import {
   BlockNoteSchema,
   combineByGroup,
-  filterSuggestionItems,
 } from "@blocknote/core";
+import { filterSuggestionItems } from "@blocknote/core/extensions";
 import * as locales from "@blocknote/core/locales";
 import "@blocknote/core/fonts/inter.css";
 import { BlockNoteView } from "@blocknote/mantine";
diff --git a/examples/03-ui-components/04-side-menu-buttons/src/App.tsx b/examples/03-ui-components/04-side-menu-buttons/src/App.tsx
index 8661c08baa..96ef099ef3 100644
--- a/examples/03-ui-components/04-side-menu-buttons/src/App.tsx
+++ b/examples/03-ui-components/04-side-menu-buttons/src/App.tsx
@@ -39,7 +39,7 @@ export default function App() {
         sideMenu={(props) => (
           <SideMenu {...props}>
             {/* Button which removes the hovered block. */}
-            <RemoveBlockButton {...props} />
+            <RemoveBlockButton />
             <DragHandleButton {...props} />
           </SideMenu>
         )}
diff --git a/examples/03-ui-components/04-side-menu-buttons/src/RemoveBlockButton.tsx b/examples/03-ui-components/04-side-menu-buttons/src/RemoveBlockButton.tsx
index c387f67f1e..774adc51b2 100644
--- a/examples/03-ui-components/04-side-menu-buttons/src/RemoveBlockButton.tsx
+++ b/examples/03-ui-components/04-side-menu-buttons/src/RemoveBlockButton.tsx
@@ -1,16 +1,20 @@
+import {} from "@blocknote/core";
+import { SideMenu } from "@blocknote/core/extensions";
 import {
-  SideMenuProps,
   useBlockNoteEditor,
   useComponentsContext,
+  useExtension,
 } from "@blocknote/react";
 import { MdDelete } from "react-icons/md";
 
 // Custom Side Menu button to remove the hovered block.
-export function RemoveBlockButton(props: SideMenuProps) {
+export function RemoveBlockButton() {
   const editor = useBlockNoteEditor();
 
   const Components = useComponentsContext()!;
 
+  const sideMenu = useExtension(SideMenu);
+
   return (
     <Components.SideMenu.Button
       label="Remove block"
@@ -18,7 +22,7 @@ export function RemoveBlockButton(props: SideMenuProps) {
         <MdDelete
           size={24}
           onClick={() => {
-            editor.removeBlocks([props.block]);
+            editor.removeBlocks([sideMenu.store.state!.block]);
           }}
         />
       }
diff --git a/examples/03-ui-components/05-side-menu-drag-handle-items/src/App.tsx b/examples/03-ui-components/05-side-menu-drag-handle-items/src/App.tsx
index 88d82aad2d..0ff3a07174 100644
--- a/examples/03-ui-components/05-side-menu-drag-handle-items/src/App.tsx
+++ b/examples/03-ui-components/05-side-menu-drag-handle-items/src/App.tsx
@@ -4,7 +4,6 @@ import "@blocknote/mantine/style.css";
 import {
   BlockColorsItem,
   DragHandleMenu,
-  DragHandleMenuProps,
   RemoveBlockItem,
   SideMenu,
   SideMenuController,
@@ -16,12 +15,12 @@ import { ResetBlockTypeItem } from "./ResetBlockTypeItem";
 // To avoid rendering issues, it's good practice to define your custom drag
 // handle menu in a separate component, instead of inline within the `sideMenu`
 // prop of `SideMenuController`.
-const CustomDragHandleMenu = (props: DragHandleMenuProps) => (
-  <DragHandleMenu {...props}>
-    <RemoveBlockItem {...props}>Delete</RemoveBlockItem>
-    <BlockColorsItem {...props}>Colors</BlockColorsItem>
+const CustomDragHandleMenu = () => (
+  <DragHandleMenu>
+    <RemoveBlockItem>Delete</RemoveBlockItem>
+    <BlockColorsItem>Colors</BlockColorsItem>
     {/* Item which resets the hovered block's type. */}
-    <ResetBlockTypeItem {...props}>Reset Type</ResetBlockTypeItem>
+    <ResetBlockTypeItem>Reset Type</ResetBlockTypeItem>
   </DragHandleMenu>
 );
 
diff --git a/examples/03-ui-components/05-side-menu-drag-handle-items/src/ResetBlockTypeItem.tsx b/examples/03-ui-components/05-side-menu-drag-handle-items/src/ResetBlockTypeItem.tsx
index 847c7b9f6c..fe446c205d 100644
--- a/examples/03-ui-components/05-side-menu-drag-handle-items/src/ResetBlockTypeItem.tsx
+++ b/examples/03-ui-components/05-side-menu-drag-handle-items/src/ResetBlockTypeItem.tsx
@@ -1,21 +1,26 @@
+import {} from "@blocknote/core";
+import { SideMenu } from "@blocknote/core/extensions";
 import {
-  DragHandleMenuProps,
   useBlockNoteEditor,
   useComponentsContext,
+  useExtension,
 } from "@blocknote/react";
+import { ReactNode } from "react";
 
-export function ResetBlockTypeItem(props: DragHandleMenuProps) {
+export function ResetBlockTypeItem(props: { children: ReactNode }) {
   const editor = useBlockNoteEditor();
 
   const Components = useComponentsContext()!;
 
+  const sideMenu = useExtension(SideMenu);
+
   return (
     <Components.Generic.Menu.Item
       onClick={() => {
-        editor.updateBlock(props.block, { type: "paragraph" });
+        editor.updateBlock(sideMenu.store.state!.block, { type: "paragraph" });
       }}
     >
-      Reset Type
+      {props.children}
     </Components.Generic.Menu.Item>
   );
 }
diff --git a/examples/03-ui-components/06-suggestion-menus-slash-menu-items/src/App.tsx b/examples/03-ui-components/06-suggestion-menus-slash-menu-items/src/App.tsx
index 61fc407ff7..c5f65cfdd8 100644
--- a/examples/03-ui-components/06-suggestion-menus-slash-menu-items/src/App.tsx
+++ b/examples/03-ui-components/06-suggestion-menus-slash-menu-items/src/App.tsx
@@ -1,8 +1,8 @@
+import { BlockNoteEditor } from "@blocknote/core";
 import {
-  BlockNoteEditor,
   filterSuggestionItems,
-  insertOrUpdateBlock,
-} from "@blocknote/core";
+  insertOrUpdateBlockForSlashMenu,
+} from "@blocknote/core/extensions";
 import "@blocknote/core/fonts/inter.css";
 import { BlockNoteView } from "@blocknote/mantine";
 import "@blocknote/mantine/style.css";
@@ -22,7 +22,7 @@ const insertHelloWorldItem = (editor: BlockNoteEditor) => ({
     // changes its type to the provided block. Otherwise, it inserts the new
     // block below and moves the text caret to it. We use this function with
     // a block containing 'Hello World' in bold.
-    insertOrUpdateBlock(editor, {
+    insertOrUpdateBlockForSlashMenu(editor, {
       type: "paragraph",
       content: [{ type: "text", text: "Hello World", styles: { bold: true } }],
     }),
diff --git a/examples/03-ui-components/10-suggestion-menus-grid-mentions/src/App.tsx b/examples/03-ui-components/10-suggestion-menus-grid-mentions/src/App.tsx
index 333c645135..0c2f9f6c21 100644
--- a/examples/03-ui-components/10-suggestion-menus-grid-mentions/src/App.tsx
+++ b/examples/03-ui-components/10-suggestion-menus-grid-mentions/src/App.tsx
@@ -1,8 +1,8 @@
 import {
   BlockNoteSchema,
   defaultInlineContentSpecs,
-  filterSuggestionItems,
 } from "@blocknote/core";
+import { filterSuggestionItems } from "@blocknote/core/extensions";
 import "@blocknote/core/fonts/inter.css";
 import { BlockNoteView } from "@blocknote/mantine";
 import "@blocknote/mantine/style.css";
diff --git a/examples/03-ui-components/11-uppy-file-panel/src/FileReplaceButton.tsx b/examples/03-ui-components/11-uppy-file-panel/src/FileReplaceButton.tsx
index e73723a9f9..6fccc130e4 100644
--- a/examples/03-ui-components/11-uppy-file-panel/src/FileReplaceButton.tsx
+++ b/examples/03-ui-components/11-uppy-file-panel/src/FileReplaceButton.tsx
@@ -70,7 +70,7 @@ export const FileReplaceButton = () => {
         variant={"panel-popover"}
       >
         {/* Replaces default file panel with our Uppy one. */}
-        <UppyFilePanel block={block as any} />
+        <UppyFilePanel blockId={block.id} />
       </Components.Generic.Popover.Content>
     </Components.Generic.Popover.Root>
   );
diff --git a/examples/03-ui-components/11-uppy-file-panel/src/UppyFilePanel.tsx b/examples/03-ui-components/11-uppy-file-panel/src/UppyFilePanel.tsx
index eaf2d4c253..4094bc4441 100644
--- a/examples/03-ui-components/11-uppy-file-panel/src/UppyFilePanel.tsx
+++ b/examples/03-ui-components/11-uppy-file-panel/src/UppyFilePanel.tsx
@@ -43,7 +43,7 @@ const uppy = new Uppy()
   });
 
 export function UppyFilePanel(props: FilePanelProps) {
-  const { block } = props;
+  const { blockId } = props;
   const editor = useBlockNoteEditor();
 
   useEffect(() => {
@@ -68,7 +68,7 @@ export function UppyFilePanel(props: FilePanelProps) {
             url: response.uploadURL,
           },
         };
-        editor.updateBlock(block, updateData);
+        editor.updateBlock(blockId, updateData);
 
         // File should be removed from the Uppy instance after upload.
         uppy.removeFile(file.id);
@@ -78,7 +78,7 @@ export function UppyFilePanel(props: FilePanelProps) {
     return () => {
       uppy.off("upload-success", handler);
     };
-  }, [block, editor]);
+  }, [blockId, editor]);
 
   // set up dashboard as in https://uppy.io/examples/
   return <Dashboard uppy={uppy} width={400} height={500} />;
diff --git a/examples/03-ui-components/13-custom-ui/src/App.tsx b/examples/03-ui-components/13-custom-ui/src/App.tsx
index 4fe6427ded..3a351264cd 100644
--- a/examples/03-ui-components/13-custom-ui/src/App.tsx
+++ b/examples/03-ui-components/13-custom-ui/src/App.tsx
@@ -1,4 +1,5 @@
-import { filterSuggestionItems } from "@blocknote/core";
+import {  } from "@blocknote/core";
+import { filterSuggestionItems } from "@blocknote/core/extensions";
 import "@blocknote/core/fonts/inter.css";
 import {
   BlockNoteViewRaw,
diff --git a/examples/03-ui-components/13-custom-ui/src/MUISideMenu.tsx b/examples/03-ui-components/13-custom-ui/src/MUISideMenu.tsx
index 8f42580c1f..e0b25b1b23 100644
--- a/examples/03-ui-components/13-custom-ui/src/MUISideMenu.tsx
+++ b/examples/03-ui-components/13-custom-ui/src/MUISideMenu.tsx
@@ -1,4 +1,11 @@
-import { SideMenuProps } from "@blocknote/react";
+import {} from "@blocknote/core";
+import { SideMenu } from "@blocknote/core/extensions";
+import {
+  SideMenuProps,
+  useBlockNoteEditor,
+  useExtension,
+  useExtensionState,
+} from "@blocknote/react";
 import { Delete, DragIndicator } from "@mui/icons-material";
 import {
   Box,
@@ -9,22 +16,22 @@ import {
   MenuItem,
   Typography,
 } from "@mui/material";
-import { MouseEvent, ReactNode, useCallback, useMemo, useState } from "react";
-
-import { TextBlockSchema } from "./schema";
+import { MouseEvent, ReactNode, useCallback, useState } from "react";
 
 // This replaces the default `RemoveBlockItem` component with a simplified
 // MUI version:
 // https://github.com/TypeCellOS/BlockNote/blob/main/packages/react/src/components/SideMenu/DragHandleMenu/DefaultItems/RemoveBlockItem.tsx
 function MUIRemoveBlockItem(
-  props: SideMenuProps<TextBlockSchema> & { closeDragHandleMenu: () => void },
+  props: SideMenuProps & { closeDragHandleMenu: () => void },
 ) {
+  const editor = useBlockNoteEditor();
+  const sideMenu = useExtension(SideMenu, { editor });
   // Deletes the block next to the side menu.
   const onClick = useCallback(() => {
-    props.unfreezeMenu();
+    sideMenu.unfreezeMenu();
     props.closeDragHandleMenu();
-    props.editor.removeBlocks([props.editor.getTextCursorPosition().block]);
-    props.editor.focus();
+    editor.removeBlocks([editor.getTextCursorPosition().block]);
+    editor.focus();
   }, [props]);
 
   return (
@@ -70,17 +77,19 @@ function MUIDragHandleMenu(props: {
 // This replaces the default `DragHandleButton` component with a simplified MUI
 // version:
 // https://github.com/TypeCellOS/BlockNote/blob/main/packages/react/src/components/SideMenu/DefaultButtons/DragHandleButton.tsx
-function MUIDragHandleButton(props: SideMenuProps<TextBlockSchema>) {
+function MUIDragHandleButton(props: SideMenuProps) {
   // Anchor/trigger element for the color menu.
   const [anchorEl, setAnchorEl] = useState<null | HTMLElement>(null);
 
+  const editor = useBlockNoteEditor();
+  const sideMenu = useExtension(SideMenu, { editor });
   // Handles opening and closing the drag handle menu.
   const onClick = useCallback(
     (event: MouseEvent<HTMLButtonElement>) => {
-      props.freezeMenu();
+      sideMenu.freezeMenu();
       setAnchorEl(event.currentTarget);
     },
-    [props],
+    [sideMenu],
   );
   const onClose = useCallback(() => {
     setAnchorEl(null);
@@ -93,8 +102,10 @@ function MUIDragHandleButton(props: SideMenuProps<TextBlockSchema>) {
         component={"button"}
         draggable={"true"}
         onClick={onClick}
-        onDragStart={(e) => props.blockDragStart(e, props.block)}
-        onDragEnd={props.blockDragEnd}
+        onDragStart={(e) =>
+          sideMenu.blockDragStart(e, sideMenu.store.state!.block)
+        }
+        onDragEnd={sideMenu.blockDragEnd}
       >
         <DragIndicator
           sx={{
@@ -115,29 +126,30 @@ function MUIDragHandleButton(props: SideMenuProps<TextBlockSchema>) {
 
 // This replaces the generic Mantine `SideMenu` component:
 // https://github.com/TypeCellOS/BlockNote/blob/main/packages/mantine/src/sideMenu/SideMenu.tsx
-function MUISideMenu(
-  props: SideMenuProps<TextBlockSchema> & { children: ReactNode },
-) {
+function MUISideMenu(props: SideMenuProps & { children: ReactNode }) {
   // Since the side menu is positioned by the top-left corner of a block, we
   // manually set its height based on the hovered block so that it's vertically
   // centered.
-  const sideMenuHeight = useMemo(() => {
-    if (props.block.type === "heading") {
-      if (props.block.props.level === 1) {
-        return 78;
-      }
+  const sideMenuHeight = useExtensionState(SideMenu, {
+    selector: (state) => {
+      // TODO this feels like a hack
+      if (state && state.block.type === "heading") {
+        if (state.block.props.level === 1) {
+          return 78;
+        }
 
-      if (props.block.props.level === 2) {
-        return 54;
-      }
+        if (state.block.props.level === 2) {
+          return 54;
+        }
 
-      if (props.block.props.level === 3) {
-        return 37;
+        if (state.block.props.level === 3) {
+          return 37;
+        }
       }
-    }
 
-    return 30;
-  }, [props.block]);
+      return 30;
+    },
+  });
 
   return (
     <Box
@@ -160,7 +172,7 @@ function MUISideMenu(
 // components specific to the Side Menu since there is really nothing more to
 // them than just an MUI `IconButton` and some styles passed via the `sx` prop,
 // as you can see in `MUIDragHandleButton`.
-export function CustomMUISideMenu(props: SideMenuProps<TextBlockSchema>) {
+export function CustomMUISideMenu(props: SideMenuProps) {
   return (
     <MUISideMenu {...props}>
       <MUIDragHandleButton {...props} />
diff --git a/examples/03-ui-components/16-link-toolbar-buttons/src/App.tsx b/examples/03-ui-components/16-link-toolbar-buttons/src/App.tsx
index 744a1e2bd1..6f28b0458c 100644
--- a/examples/03-ui-components/16-link-toolbar-buttons/src/App.tsx
+++ b/examples/03-ui-components/16-link-toolbar-buttons/src/App.tsx
@@ -55,10 +55,10 @@ export default function App() {
             <EditLinkButton
               url={props.url}
               text={props.text}
-              editLink={props.editLink}
+              range={props.range}
             />
             <OpenLinkButton url={props.url} />
-            <DeleteLinkButton deleteLink={props.deleteLink} />
+            <DeleteLinkButton />
             {/* Extra button to open alert. */}
             <AlertButton {...props} />
           </LinkToolbar>
diff --git a/examples/05-interoperability/05-converting-blocks-to-pdf/src/App.tsx b/examples/05-interoperability/05-converting-blocks-to-pdf/src/App.tsx
index d910886113..7a34932ec8 100644
--- a/examples/05-interoperability/05-converting-blocks-to-pdf/src/App.tsx
+++ b/examples/05-interoperability/05-converting-blocks-to-pdf/src/App.tsx
@@ -1,9 +1,9 @@
 import {
   BlockNoteSchema,
   combineByGroup,
-  filterSuggestionItems,
   withPageBreak,
 } from "@blocknote/core";
+import { filterSuggestionItems } from "@blocknote/core/extensions";
 import "@blocknote/core/fonts/inter.css";
 import * as locales from "@blocknote/core/locales";
 import { BlockNoteView } from "@blocknote/mantine";
diff --git a/examples/05-interoperability/06-converting-blocks-to-docx/src/App.tsx b/examples/05-interoperability/06-converting-blocks-to-docx/src/App.tsx
index 9c91a99aff..8c5ae2398f 100644
--- a/examples/05-interoperability/06-converting-blocks-to-docx/src/App.tsx
+++ b/examples/05-interoperability/06-converting-blocks-to-docx/src/App.tsx
@@ -1,9 +1,9 @@
 import {
   BlockNoteSchema,
   combineByGroup,
-  filterSuggestionItems,
   withPageBreak,
 } from "@blocknote/core";
+import { filterSuggestionItems } from "@blocknote/core/extensions";
 import "@blocknote/core/fonts/inter.css";
 import * as locales from "@blocknote/core/locales";
 import { BlockNoteView } from "@blocknote/mantine";
diff --git a/examples/05-interoperability/07-converting-blocks-to-odt/src/App.tsx b/examples/05-interoperability/07-converting-blocks-to-odt/src/App.tsx
index 9ece22d905..fcfa07ff85 100644
--- a/examples/05-interoperability/07-converting-blocks-to-odt/src/App.tsx
+++ b/examples/05-interoperability/07-converting-blocks-to-odt/src/App.tsx
@@ -2,8 +2,8 @@ import {
   BlockNoteSchema,
   combineByGroup,
   withPageBreak,
-  filterSuggestionItems,
 } from "@blocknote/core";
+import { filterSuggestionItems } from "@blocknote/core/extensions";
 import * as locales from "@blocknote/core/locales";
 import "@blocknote/core/fonts/inter.css";
 import { BlockNoteView } from "@blocknote/mantine";
diff --git a/examples/05-interoperability/08-converting-blocks-to-react-email/src/App.tsx b/examples/05-interoperability/08-converting-blocks-to-react-email/src/App.tsx
index 9b46d8e927..5d3d896b8f 100644
--- a/examples/05-interoperability/08-converting-blocks-to-react-email/src/App.tsx
+++ b/examples/05-interoperability/08-converting-blocks-to-react-email/src/App.tsx
@@ -3,9 +3,9 @@ import {
   COLORS_DARK_MODE_DEFAULT,
   COLORS_DEFAULT,
   combineByGroup,
-  filterSuggestionItems,
   withPageBreak,
 } from "@blocknote/core";
+import { filterSuggestionItems } from "@blocknote/core/extensions";
 import "@blocknote/core/fonts/inter.css";
 import { BlockNoteView } from "@blocknote/mantine";
 import "@blocknote/mantine/style.css";
diff --git a/examples/06-custom-schema/02-suggestion-menus-mentions/src/App.tsx b/examples/06-custom-schema/02-suggestion-menus-mentions/src/App.tsx
index 82bbccdf59..4339153441 100644
--- a/examples/06-custom-schema/02-suggestion-menus-mentions/src/App.tsx
+++ b/examples/06-custom-schema/02-suggestion-menus-mentions/src/App.tsx
@@ -1,8 +1,8 @@
 import {
   BlockNoteSchema,
   defaultInlineContentSpecs,
-  filterSuggestionItems,
 } from "@blocknote/core";
+import { filterSuggestionItems } from "@blocknote/core/extensions";
 import "@blocknote/core/fonts/inter.css";
 import { BlockNoteView } from "@blocknote/mantine";
 import "@blocknote/mantine/style.css";
diff --git a/examples/06-custom-schema/04-pdf-file-block/src/App.tsx b/examples/06-custom-schema/04-pdf-file-block/src/App.tsx
index f0252e4dc4..3c244a2840 100644
--- a/examples/06-custom-schema/04-pdf-file-block/src/App.tsx
+++ b/examples/06-custom-schema/04-pdf-file-block/src/App.tsx
@@ -1,9 +1,8 @@
+import { BlockNoteSchema, defaultBlockSpecs } from "@blocknote/core";
 import {
-  BlockNoteSchema,
-  defaultBlockSpecs,
   filterSuggestionItems,
-  insertOrUpdateBlock,
-} from "@blocknote/core";
+  insertOrUpdateBlockForSlashMenu,
+} from "@blocknote/core/extensions";
 import { en } from "@blocknote/core/locales";
 import "@blocknote/core/fonts/inter.css";
 import { BlockNoteView } from "@blocknote/mantine";
@@ -39,7 +38,7 @@ const schema = BlockNoteSchema.create({
 const insertPDF = (editor: typeof schema.BlockNoteEditor) => ({
   title: "PDF",
   onItemClick: () => {
-    insertOrUpdateBlock(editor, {
+    insertOrUpdateBlockForSlashMenu(editor, {
       type: "pdf",
     });
   },
diff --git a/examples/06-custom-schema/05-alert-block-full-ux/src/App.tsx b/examples/06-custom-schema/05-alert-block-full-ux/src/App.tsx
index 21dc29650a..625dcce896 100644
--- a/examples/06-custom-schema/05-alert-block-full-ux/src/App.tsx
+++ b/examples/06-custom-schema/05-alert-block-full-ux/src/App.tsx
@@ -1,9 +1,8 @@
+import { BlockNoteSchema, defaultBlockSpecs } from "@blocknote/core";
 import {
-  BlockNoteSchema,
-  defaultBlockSpecs,
   filterSuggestionItems,
-  insertOrUpdateBlock,
-} from "@blocknote/core";
+  insertOrUpdateBlockForSlashMenu,
+} from "@blocknote/core/extensions";
 import "@blocknote/core/fonts/inter.css";
 import { BlockNoteView } from "@blocknote/mantine";
 import "@blocknote/mantine/style.css";
@@ -38,7 +37,7 @@ const insertAlert = (editor: typeof schema.BlockNoteEditor) => ({
     // changes its type to the provided block. Otherwise, it inserts the new
     // block below and moves the text caret to it. We use this function with an
     // Alert block.
-    insertOrUpdateBlock(editor, {
+    insertOrUpdateBlockForSlashMenu(editor, {
       type: "alert",
     }),
   aliases: [
diff --git a/examples/07-collaboration/08-forking/src/App.tsx b/examples/07-collaboration/08-forking/src/App.tsx
index ba6398bb51..134935d546 100644
--- a/examples/07-collaboration/08-forking/src/App.tsx
+++ b/examples/07-collaboration/08-forking/src/App.tsx
@@ -1,11 +1,15 @@
 import "@blocknote/core/fonts/inter.css";
-import { useCreateBlockNote } from "@blocknote/react";
+import {} from "@blocknote/core";
+import { ForkYDoc } from "@blocknote/core/extensions";
+import {
+  useCreateBlockNote,
+  useExtension,
+  useExtensionState,
+} from "@blocknote/react";
 import { BlockNoteView } from "@blocknote/mantine";
 import "@blocknote/mantine/style.css";
 import YPartyKitProvider from "y-partykit/provider";
 import * as Y from "yjs";
-import { useEffect } from "react";
-import { useState } from "react";
 
 // Sets up Yjs document and PartyKit Yjs provider.
 const doc = new Y.Doc();
@@ -30,18 +34,18 @@ export default function App() {
       },
     },
   });
-  const [isForked, setIsForked] = useState(false);
-
-  useEffect(() => {
-    editor.forkYDocPlugin!.on("forked", setIsForked);
-  }, [editor]);
+  const forkYDocPlugin = useExtension(ForkYDoc, { editor });
+  const isForked = useExtensionState(ForkYDoc, {
+    editor,
+    selector: (state) => state.isForked,
+  });
 
   // Renders the editor instance.
   return (
     <>
       <button
         onClick={() => {
-          editor.forkYDocPlugin!.fork();
+          forkYDocPlugin.fork();
         }}
         disabled={isForked}
       >
@@ -49,7 +53,7 @@ export default function App() {
       </button>
       <button
         onClick={() => {
-          editor.forkYDocPlugin!.merge({ keepChanges: true });
+          forkYDocPlugin.merge({ keepChanges: true });
         }}
         disabled={!isForked}
       >
@@ -57,7 +61,7 @@ export default function App() {
       </button>
       <button
         onClick={() => {
-          editor.forkYDocPlugin!.merge({ keepChanges: false });
+          forkYDocPlugin.merge({ keepChanges: false });
         }}
         disabled={!isForked}
       >
diff --git a/examples/09-ai/01-minimal/src/App.tsx b/examples/09-ai/01-minimal/src/App.tsx
index 8c0ce912a2..f5cc569f39 100644
--- a/examples/09-ai/01-minimal/src/App.tsx
+++ b/examples/09-ai/01-minimal/src/App.tsx
@@ -1,4 +1,5 @@
-import { BlockNoteEditor, filterSuggestionItems } from "@blocknote/core";
+import { BlockNoteEditor,  } from "@blocknote/core";
+import { filterSuggestionItems } from "@blocknote/core/extensions";
 import "@blocknote/core/fonts/inter.css";
 import { en } from "@blocknote/core/locales";
 import { BlockNoteView } from "@blocknote/mantine";
@@ -12,9 +13,9 @@ import {
   useCreateBlockNote,
 } from "@blocknote/react";
 import {
+  AIExtension,
   AIMenuController,
   AIToolbarButton,
-  createAIExtension,
   getAISlashMenuItems,
 } from "@blocknote/xl-ai";
 import { en as aiEn } from "@blocknote/xl-ai/locales";
@@ -34,14 +35,13 @@ export default function App() {
       ai: aiEn, // add default translations for the AI extension
     },
     // Register the AI extension
-    extensions: [
-      createAIExtension({
-        transport: new DefaultChatTransport({
-          // URL to your backend API, see example source in `packages/xl-ai-server/src/routes/regular.ts`
-          api: `${BASE_URL}/regular/streamText`,
-        }),
+    extensions: [AIExtension],
+    ai: {
+      transport: new DefaultChatTransport({
+        // URL to your backend API, see example source in `packages/xl-ai-server/src/routes/regular.ts`
+        api: `${BASE_URL}/regular/streamText`,
       }),
-    ],
+    },
     // We set some initial content for demo purposes
     initialContent: [
       {
diff --git a/examples/09-ai/02-playground/src/App.tsx b/examples/09-ai/02-playground/src/App.tsx
index 17ceacd057..76f6cb938a 100644
--- a/examples/09-ai/02-playground/src/App.tsx
+++ b/examples/09-ai/02-playground/src/App.tsx
@@ -1,4 +1,5 @@
-import { BlockNoteEditor, filterSuggestionItems } from "@blocknote/core";
+import { BlockNoteEditor } from "@blocknote/core";
+import { filterSuggestionItems } from "@blocknote/core/extensions";
 import "@blocknote/core/fonts/inter.css";
 import { en } from "@blocknote/core/locales";
 import { BlockNoteView } from "@blocknote/mantine";
@@ -11,13 +12,13 @@ import {
   getFormattingToolbarItems,
   useBlockNoteContext,
   useCreateBlockNote,
+  useExtension,
   usePrefersColorScheme,
 } from "@blocknote/react";
 import {
+  AIExtension,
   AIMenuController,
   AIToolbarButton,
-  createAIExtension,
-  getAIExtension,
   getAISlashMenuItems,
 } from "@blocknote/xl-ai";
 import { en as aiEn } from "@blocknote/xl-ai/locales";
@@ -44,14 +45,13 @@ export default function App() {
       ai: aiEn, // add default translations for the AI extension
     },
     // Register the AI extension
-    extensions: [
-      createAIExtension({
-        transport: new DefaultChatTransport({
-          // URL to your backend API, see example source in `packages/xl-ai-server/src/routes/regular.ts`
-          api: `${BASE_URL}/model-playground/streamText`,
-        }),
+    extensions: [AIExtension],
+    ai: {
+      transport: new DefaultChatTransport({
+        // URL to your backend API, see example source in `packages/xl-ai-server/src/routes/regular.ts`
+        api: `${BASE_URL}/model-playground/streamText`,
       }),
-    ],
+    },
     // We set some initial content for demo purposes
     initialContent: [
       {
@@ -79,7 +79,7 @@ export default function App() {
     ],
   });
 
-  const ai = getAIExtension(editor);
+  const ai = useExtension(AIExtension, { editor });
 
   useEffect(() => {
     // update the default model in the extension
@@ -92,7 +92,7 @@ export default function App() {
         },
       },
     });
-  }, [model, ai.options]);
+  }, [model, ai]);
 
   const themePreference = usePrefersColorScheme();
   const existingContext = useBlockNoteContext();
diff --git a/examples/09-ai/03-custom-ai-menu-items/src/App.tsx b/examples/09-ai/03-custom-ai-menu-items/src/App.tsx
index 876e371b4d..e12907c46d 100644
--- a/examples/09-ai/03-custom-ai-menu-items/src/App.tsx
+++ b/examples/09-ai/03-custom-ai-menu-items/src/App.tsx
@@ -1,4 +1,5 @@
-import { BlockNoteEditor, filterSuggestionItems } from "@blocknote/core";
+import { BlockNoteEditor,  } from "@blocknote/core";
+import { filterSuggestionItems } from "@blocknote/core/extensions";
 import "@blocknote/core/fonts/inter.css";
 import { en } from "@blocknote/core/locales";
 import { BlockNoteView } from "@blocknote/mantine";
@@ -12,10 +13,10 @@ import {
   useCreateBlockNote,
 } from "@blocknote/react";
 import {
+  AIExtension,
   AIMenu,
   AIMenuController,
   AIToolbarButton,
-  createAIExtension,
   getAISlashMenuItems,
   getDefaultAIMenuItems,
 } from "@blocknote/xl-ai";
@@ -37,13 +38,12 @@ export default function App() {
       ai: aiEn, // add default translations for the AI extension
     },
     // Register the AI extension
-    extensions: [
-      createAIExtension({
-        transport: new DefaultChatTransport({
-          api: `${BASE_URL}/regular/streamText`,
-        }),
+    extensions: [AIExtension],
+    ai: {
+      transport: new DefaultChatTransport({
+        api: `${BASE_URL}/regular/streamText`,
       }),
-    ],
+    },
     // We set some initial content for demo purposes
     initialContent: [
       {
diff --git a/examples/09-ai/03-custom-ai-menu-items/src/customAIMenuItems.tsx b/examples/09-ai/03-custom-ai-menu-items/src/customAIMenuItems.tsx
index 17a4f75bfe..526589cdf2 100644
--- a/examples/09-ai/03-custom-ai-menu-items/src/customAIMenuItems.tsx
+++ b/examples/09-ai/03-custom-ai-menu-items/src/customAIMenuItems.tsx
@@ -1,8 +1,8 @@
 import { BlockNoteEditor } from "@blocknote/core";
 import {
+  AIExtension,
   aiDocumentFormats,
   AIMenuSuggestionItem,
-  getAIExtension,
 } from "@blocknote/xl-ai";
 import { RiApps2AddFill, RiEmotionHappyFill } from "react-icons/ri";
 
@@ -17,7 +17,7 @@ export const makeInformal = (
   aliases: ["informal", "make informal", "casual"],
   icon: <RiEmotionHappyFill size={18} />,
   onItemClick: async () => {
-    await getAIExtension(editor).invokeAI({
+    await editor.getExtension(AIExtension)?.invokeAI({
       userPrompt: "Give the selected text a more informal (casual) tone",
       // Set to true to tell the LLM to specifically
       // use the selected content as context. Defaults
@@ -51,7 +51,7 @@ export const addRelatedTopics = (
   aliases: ["related topics", "find topics"],
   icon: <RiApps2AddFill size={18} />,
   onItemClick: async () => {
-    await getAIExtension(editor).invokeAI({
+    await editor.getExtension(AIExtension)?.invokeAI({
       userPrompt:
         "Think of some related topics to the current text and write a sentence about each",
       // Sets what operations the LLM is allowed to do.
diff --git a/examples/09-ai/04-with-collaboration/src/App.tsx b/examples/09-ai/04-with-collaboration/src/App.tsx
index 6ffbf301d4..97f3b06aa0 100644
--- a/examples/09-ai/04-with-collaboration/src/App.tsx
+++ b/examples/09-ai/04-with-collaboration/src/App.tsx
@@ -1,4 +1,5 @@
-import { BlockNoteEditor, filterSuggestionItems } from "@blocknote/core";
+import { BlockNoteEditor,  } from "@blocknote/core";
+import { filterSuggestionItems } from "@blocknote/core/extensions";
 import "@blocknote/core/fonts/inter.css";
 import { en } from "@blocknote/core/locales";
 import { BlockNoteView } from "@blocknote/mantine";
@@ -12,9 +13,9 @@ import {
   useCreateBlockNote,
 } from "@blocknote/react";
 import {
+  AIExtension,
   AIMenuController,
   AIToolbarButton,
-  createAIExtension,
   getAISlashMenuItems,
 } from "@blocknote/xl-ai";
 import { en as aiEn } from "@blocknote/xl-ai/locales";
@@ -79,13 +80,12 @@ export default function App() {
       ai: aiEn, // add default translations for the AI extension
     },
     // Register the AI extension
-    extensions: [
-      createAIExtension({
-        transport: new DefaultChatTransport({
-          api: `${BASE_URL}/regular/streamText`,
-        }),
+    extensions: [AIExtension],
+    ai: {
+      transport: new DefaultChatTransport({
+        api: `${BASE_URL}/regular/streamText`,
       }),
-    ],
+    },
     // We set some initial content for demo purposes
     initialContent: [
       {
diff --git a/examples/09-ai/05-manual-execution/src/App.tsx b/examples/09-ai/05-manual-execution/src/App.tsx
index cec4f959a0..fa7d1236bc 100644
--- a/examples/09-ai/05-manual-execution/src/App.tsx
+++ b/examples/09-ai/05-manual-execution/src/App.tsx
@@ -4,10 +4,9 @@ import { BlockNoteView } from "@blocknote/mantine";
 import "@blocknote/mantine/style.css";
 import { useCreateBlockNote } from "@blocknote/react";
 import {
+  AIExtension,
   StreamToolExecutor,
   aiDocumentFormats,
-  createAIExtension,
-  getAIExtension,
 } from "@blocknote/xl-ai";
 import { en as aiEn } from "@blocknote/xl-ai/locales";
 import "@blocknote/xl-ai/style.css";
@@ -20,7 +19,7 @@ export default function App() {
       ai: aiEn, // add default translations for the AI extension
     },
     // Register the AI extension
-    extensions: [createAIExtension({})],
+    extensions: [AIExtension],
     // We set some initial content for demo purposes
     initialContent: [
       {
@@ -77,7 +76,7 @@ export default function App() {
             });
             // accept the changes after 1 second
             await new Promise((resolve) => setTimeout(resolve, 1000));
-            await getAIExtension(editor).acceptChanges();
+            await editor.getExtension(AIExtension)?.acceptChanges();
           }}
         >
           Update first block
@@ -139,7 +138,7 @@ export default function App() {
 
             // accept the changes after 1 second
             await new Promise((resolve) => setTimeout(resolve, 1000));
-            await getAIExtension(editor).acceptChanges();
+            await editor.getExtension(AIExtension)?.acceptChanges();
           }}
         >
           Update first block (streaming)
@@ -187,7 +186,7 @@ export default function App() {
 
             // accept the changes after 1 second
             await new Promise((resolve) => setTimeout(resolve, 1000));
-            await getAIExtension(editor).acceptChanges();
+            await editor.getExtension(AIExtension)?.acceptChanges();
           }}
         >
           Update first block (streaming strings)
diff --git a/examples/09-ai/06-client-side-transport/src/App.tsx b/examples/09-ai/06-client-side-transport/src/App.tsx
index b21efdadc6..a74dff390b 100644
--- a/examples/09-ai/06-client-side-transport/src/App.tsx
+++ b/examples/09-ai/06-client-side-transport/src/App.tsx
@@ -1,5 +1,6 @@
 import { createGroq } from "@ai-sdk/groq";
-import { BlockNoteEditor, filterSuggestionItems } from "@blocknote/core";
+import { BlockNoteEditor,  } from "@blocknote/core";
+import { filterSuggestionItems } from "@blocknote/core/extensions";
 import "@blocknote/core/fonts/inter.css";
 import { en } from "@blocknote/core/locales";
 import { BlockNoteView } from "@blocknote/mantine";
@@ -13,10 +14,10 @@ import {
   useCreateBlockNote,
 } from "@blocknote/react";
 import {
+  AIExtension,
   AIMenuController,
   AIToolbarButton,
   ClientSideTransport,
-  createAIExtension,
   fetchViaProxy,
   getAISlashMenuItems,
 } from "@blocknote/xl-ai";
@@ -47,16 +48,15 @@ export default function App() {
       ai: aiEn, // add default translations for the AI extension
     },
     // Register the AI extension
-    extensions: [
-      createAIExtension({
-        // The ClientSideTransport is used so the client makes calls directly to `streamText`
-        // (whereas normally in the Vercel AI SDK, the client makes calls to your server, which then calls these methods)
-        // (see https://github.com/vercel/ai/issues/5140 for background info)
-        transport: new ClientSideTransport({
-          model,
-        }),
+    extensions: [AIExtension],
+    ai: {
+      // The ClientSideTransport is used so the client makes calls directly to `streamText`
+      // (whereas normally in the Vercel AI SDK, the client makes calls to your server, which then calls these methods)
+      // (see https://github.com/vercel/ai/issues/5140 for background info)
+      transport: new ClientSideTransport({
+        model,
       }),
-    ],
+    },
     // We set some initial content for demo purposes
     initialContent: [
       {
diff --git a/examples/09-ai/07-server-promptbuilder/src/App.tsx b/examples/09-ai/07-server-promptbuilder/src/App.tsx
index 90d89e0c96..c11ec8a9fe 100644
--- a/examples/09-ai/07-server-promptbuilder/src/App.tsx
+++ b/examples/09-ai/07-server-promptbuilder/src/App.tsx
@@ -1,4 +1,5 @@
-import { BlockNoteEditor, filterSuggestionItems } from "@blocknote/core";
+import { BlockNoteEditor,  } from "@blocknote/core";
+import { filterSuggestionItems } from "@blocknote/core/extensions";
 import "@blocknote/core/fonts/inter.css";
 import { en } from "@blocknote/core/locales";
 import { BlockNoteView } from "@blocknote/mantine";
@@ -12,10 +13,10 @@ import {
   useCreateBlockNote,
 } from "@blocknote/react";
 import {
+  AIExtension,
   aiDocumentFormats,
   AIMenuController,
   AIToolbarButton,
-  createAIExtension,
   defaultAIRequestSender,
   getAISlashMenuItems,
 } from "@blocknote/xl-ai";
@@ -35,51 +36,50 @@ export default function App() {
       ai: aiEn, // add default translations for the AI extension
     },
     // Register the AI extension
-    extensions: [
-      createAIExtension({
-        // similar to https://ai-sdk.dev/docs/ai-sdk-ui/chatbot-message-persistence#sending-only-the-last-message
-        // we adjust the transport to not send all messages to the backend
-        transport: new DefaultChatTransport({
-          // (see packages/xl-ai-server/src/routes/vercelAiSdkPersistence.ts)
-          api: `${BASE_URL}/server-promptbuilder/streamText`,
-          prepareSendMessagesRequest({ id, body, messages, requestMetadata }) {
-            // we don't send the messages, just the information we need to compose / append messages server-side:
-            // - the conversation id
-            // - the promptData
-            // - the tool results of the last message
+    extensions: [AIExtension],
+    ai: {
+      // similar to https://ai-sdk.dev/docs/ai-sdk-ui/chatbot-message-persistence#sending-only-the-last-message
+      // we adjust the transport to not send all messages to the backend
+      transport: new DefaultChatTransport({
+        // (see packages/xl-ai-server/src/routes/vercelAiSdkPersistence.ts)
+        api: `${BASE_URL}/server-promptbuilder/streamText`,
+        prepareSendMessagesRequest({ id, body, messages, requestMetadata }) {
+          // we don't send the messages, just the information we need to compose / append messages server-side:
+          // - the conversation id
+          // - the promptData
+          // - the tool results of the last message
 
-            // we need to share data about tool calls with the backend,
-            // as these can be client-side executed. The backend needs to know the tool outputs
-            // in order to compose a new valid LLM request
-            const lastToolParts =
-              messages.length > 0
-                ? messages[messages.length - 1].parts.filter((part) =>
-                    isToolOrDynamicToolUIPart(part),
-                  )
-                : [];
+          // we need to share data about tool calls with the backend,
+          // as these can be client-side executed. The backend needs to know the tool outputs
+          // in order to compose a new valid LLM request
+          const lastToolParts =
+            messages.length > 0
+              ? messages[messages.length - 1].parts.filter((part) =>
+                  isToolOrDynamicToolUIPart(part),
+                )
+              : [];
 
-            return {
-              body: {
-                ...body,
-                // TODO: this conversation id is client-side generated, we
-                // should have a server-side generated id to ensure uniqueness
-                // see https://github.com/vercel/ai/issues/7340#issuecomment-3307559636
-                id,
-                // get the promptData from requestMetadata (set by `promptAIRequestSender`) and send to backend
-                promptData: (requestMetadata as any).promptData,
-                lastToolParts,
-                // messages, -> we explicitly don't send the messages array as we compose messages server-side
-              },
-            };
-          },
-        }),
-        // customize the aiRequestSender to not update the messages array on the client-side
-        aiRequestSender: defaultAIRequestSender(
-          async () => {}, // disable the client-side promptbuilder
-          aiDocumentFormats.html.defaultPromptInputDataBuilder,
-        ),
+          return {
+            body: {
+              ...body,
+              // TODO: this conversation id is client-side generated, we
+              // should have a server-side generated id to ensure uniqueness
+              // see https://github.com/vercel/ai/issues/7340#issuecomment-3307559636
+              id,
+              // get the promptData from requestMetadata (set by `promptAIRequestSender`) and send to backend
+              promptData: (requestMetadata as any).promptData,
+              lastToolParts,
+              // messages, -> we explicitly don't send the messages array as we compose messages server-side
+            },
+          };
+        },
       }),
-    ],
+      // customize the aiRequestSender to not update the messages array on the client-side
+      aiRequestSender: defaultAIRequestSender(
+        async () => {}, // disable the client-side promptbuilder
+        aiDocumentFormats.html.defaultPromptInputDataBuilder,
+      ),
+    },
     // We set some initial content for demo purposes
     initialContent: [
       {
diff --git a/packages/core/package.json b/packages/core/package.json
index 9a434e1d89..78343f8ce2 100644
--- a/packages/core/package.json
+++ b/packages/core/package.json
@@ -67,6 +67,11 @@
       "import": "./dist/locales.js",
       "require": "./dist/locales.cjs"
     },
+    "./extensions": {
+      "types": "./types/src/extensions/index.d.ts",
+      "import": "./dist/extensions.js",
+      "require": "./dist/extensions.cjs"
+    },
     "./yjs": {
       "types": "./types/src/yjs/index.d.ts",
       "import": "./dist/yjs.js",
@@ -85,12 +90,13 @@
   },
   "dependencies": {
     "@emoji-mart/data": "^1.2.1",
+    "@handlewithcare/prosemirror-inputrules": "0.1.3",
     "@shikijs/types": "3.13.0",
+    "@tanstack/store": "0.7.7",
     "@tiptap/core": "^3.11.0",
     "@tiptap/extension-bold": "^3.7.2",
     "@tiptap/extension-code": "^3.7.2",
     "@tiptap/extension-gapcursor": "^3.7.2",
-    "@tiptap/extension-history": "^3.7.2",
     "@tiptap/extension-horizontal-rule": "^3.7.2",
     "@tiptap/extension-italic": "^3.7.2",
     "@tiptap/extension-link": "^3.7.2",
diff --git a/packages/core/src/api/nodeConversions/blockToNode.ts b/packages/core/src/api/nodeConversions/blockToNode.ts
index 660d97406e..206ff8d9fd 100644
--- a/packages/core/src/api/nodeConversions/blockToNode.ts
+++ b/packages/core/src/api/nodeConversions/blockToNode.ts
@@ -1,6 +1,6 @@
 import { Attrs, Fragment, Mark, Node, Schema } from "@tiptap/pm/model";
 
-import UniqueID from "../../extensions/UniqueID/UniqueID.js";
+import UniqueID from "../../extensions/tiptap-extensions/UniqueID/UniqueID.js";
 import type {
   InlineContentSchema,
   PartialCustomInlineContentFromConfig,
diff --git a/packages/core/src/api/nodeConversions/nodeToBlock.ts b/packages/core/src/api/nodeConversions/nodeToBlock.ts
index 7ce9d39cf1..5048f91a2b 100644
--- a/packages/core/src/api/nodeConversions/nodeToBlock.ts
+++ b/packages/core/src/api/nodeConversions/nodeToBlock.ts
@@ -1,6 +1,6 @@
 import { Mark, Node, Schema, Slice } from "@tiptap/pm/model";
 import type { Block } from "../../blocks/defaultBlocks.js";
-import UniqueID from "../../extensions/UniqueID/UniqueID.js";
+import UniqueID from "../../extensions/tiptap-extensions/UniqueID/UniqueID.js";
 import type {
   BlockSchema,
   CustomInlineContentConfig,
diff --git a/packages/core/src/blocks/Code/block.ts b/packages/core/src/blocks/Code/block.ts
index 1458257e48..1886fd5afe 100644
--- a/packages/core/src/blocks/Code/block.ts
+++ b/packages/core/src/blocks/Code/block.ts
@@ -1,5 +1,5 @@
 import type { HighlighterGeneric } from "@shikijs/types";
-import { createBlockNoteExtension } from "../../editor/BlockNoteExtension.js";
+import { createExtension } from "../../editor/BlockNoteExtension.js";
 import { createBlockConfig, createBlockSpec } from "../../schema/index.js";
 import { lazyShikiPlugin } from "./shiki.js";
 import { DOMParser } from "@tiptap/pm/model";
@@ -169,11 +169,11 @@ export const createCodeBlockSpec = createBlockSpec(
   }),
   (options) => {
     return [
-      createBlockNoteExtension({
+      createExtension({
         key: "code-block-highlighter",
-        plugins: [lazyShikiPlugin(options)],
+        prosemirrorPlugins: [lazyShikiPlugin(options)],
       }),
-      createBlockNoteExtension({
+      createExtension({
         key: "code-block-keyboard-shortcuts",
         keyboardShortcuts: {
           Delete: ({ editor }) => {
diff --git a/packages/core/src/blocks/Divider/block.ts b/packages/core/src/blocks/Divider/block.ts
index 3de2211d3f..ad847ddd0c 100644
--- a/packages/core/src/blocks/Divider/block.ts
+++ b/packages/core/src/blocks/Divider/block.ts
@@ -1,4 +1,4 @@
-import { createBlockNoteExtension } from "../../editor/BlockNoteExtension.js";
+import { createExtension } from "../../editor/BlockNoteExtension.js";
 import { createBlockConfig, createBlockSpec } from "../../schema/index.js";
 
 export type DividerBlockConfig = ReturnType<typeof createDividerBlockConfig>;
@@ -34,7 +34,7 @@ export const createDividerBlockSpec = createBlockSpec(
     },
   },
   [
-    createBlockNoteExtension({
+    createExtension({
       key: "divider-block-shortcuts",
       inputRules: [
         {
diff --git a/packages/core/src/blocks/File/helpers/render/createAddFileButton.ts b/packages/core/src/blocks/File/helpers/render/createAddFileButton.ts
index 227856b4ac..8b2fbd06c7 100644
--- a/packages/core/src/blocks/File/helpers/render/createAddFileButton.ts
+++ b/packages/core/src/blocks/File/helpers/render/createAddFileButton.ts
@@ -1,4 +1,5 @@
 import type { BlockNoteEditor } from "../../../../editor/BlockNoteEditor.js";
+import { FilePanel } from "../../../../extensions/FilePanel/FilePanel.js";
 import {
   BlockConfig,
   BlockFromConfigNoChildren,
@@ -33,14 +34,11 @@ export const createAddFileButton = (
   // Prevents focus from moving to the button.
   const addFileButtonMouseDownHandler = (event: MouseEvent) => {
     event.preventDefault();
+    event.stopPropagation();
   };
   // Opens the file toolbar.
   const addFileButtonClickHandler = () => {
-    editor.transact((tr) =>
-      tr.setMeta(editor.filePanel!.plugins[0], {
-        block: block,
-      }),
-    );
+    editor.getExtension(FilePanel)?.showMenu(block.id);
   };
   addFileButton.addEventListener(
     "mousedown",
diff --git a/packages/core/src/blocks/Heading/block.ts b/packages/core/src/blocks/Heading/block.ts
index 7181298d1f..c5727b8ec9 100644
--- a/packages/core/src/blocks/Heading/block.ts
+++ b/packages/core/src/blocks/Heading/block.ts
@@ -1,5 +1,6 @@
 import { createBlockConfig, createBlockSpec } from "../../schema/index.js";
-import { createBlockNoteExtension } from "../../editor/BlockNoteExtension.js";
+import { BlockNoteEditor } from "../../editor/BlockNoteEditor.js";
+import { createExtension } from "../../editor/BlockNoteExtension.js";
 import {
   addDefaultPropsExternalHTML,
   defaultProps,
@@ -16,6 +17,25 @@ export interface HeadingOptions {
   allowToggleHeadings?: boolean;
 }
 
+const createHeadingKeyboardShortcut =
+  (level: number) =>
+  ({ editor }: { editor: BlockNoteEditor<any, any, any> }) => {
+    const cursorPosition = editor.getTextCursorPosition();
+
+    if (
+      editor.schema.blockSchema[cursorPosition.block.type].content !== "inline"
+    ) {
+      return false;
+    }
+
+    editor.updateBlock(cursorPosition.block, {
+      type: "heading",
+      props: { level },
+    });
+
+    return true;
+  };
+
 export type HeadingBlockConfig = ReturnType<typeof createHeadingBlockConfig>;
 
 export const createHeadingBlockConfig = createBlockConfig(
@@ -97,29 +117,12 @@ export const createHeadingBlockSpec = createBlockSpec(
     },
   }),
   ({ levels = HEADING_LEVELS }: HeadingOptions = {}) => [
-    createBlockNoteExtension({
+    createExtension({
       key: "heading-shortcuts",
       keyboardShortcuts: Object.fromEntries(
         levels.map((level) => [
           `Mod-Alt-${level}`,
-          ({ editor }) => {
-            const cursorPosition = editor.getTextCursorPosition();
-
-            if (
-              editor.schema.blockSchema[cursorPosition.block.type].content !==
-              "inline"
-            ) {
-              return false;
-            }
-
-            editor.updateBlock(cursorPosition.block, {
-              type: "heading",
-              props: {
-                level: level as any,
-              },
-            });
-            return true;
-          },
+          createHeadingKeyboardShortcut(level),
         ]) ?? [],
       ),
       inputRules: levels.map((level) => ({
diff --git a/packages/core/src/blocks/ListItem/BulletListItem/block.ts b/packages/core/src/blocks/ListItem/BulletListItem/block.ts
index 029538d4a3..2314d09ced 100644
--- a/packages/core/src/blocks/ListItem/BulletListItem/block.ts
+++ b/packages/core/src/blocks/ListItem/BulletListItem/block.ts
@@ -1,5 +1,5 @@
 import { getBlockInfoFromSelection } from "../../../api/getBlockInfoFromPos.js";
-import { createBlockNoteExtension } from "../../../editor/BlockNoteExtension.js";
+import { createExtension } from "../../../editor/BlockNoteExtension.js";
 import { createBlockConfig, createBlockSpec } from "../../../schema/index.js";
 import {
   addDefaultPropsExternalHTML,
@@ -78,7 +78,7 @@ export const createBulletListItemBlockSpec = createBlockSpec(
     },
   },
   [
-    createBlockNoteExtension({
+    createExtension({
       key: "bullet-list-item-shortcuts",
       keyboardShortcuts: {
         Enter: ({ editor }) => {
diff --git a/packages/core/src/blocks/ListItem/CheckListItem/block.ts b/packages/core/src/blocks/ListItem/CheckListItem/block.ts
index 7887144b1b..f954e8a429 100644
--- a/packages/core/src/blocks/ListItem/CheckListItem/block.ts
+++ b/packages/core/src/blocks/ListItem/CheckListItem/block.ts
@@ -1,4 +1,4 @@
-import { createBlockNoteExtension } from "../../../editor/BlockNoteExtension.js";
+import { createExtension } from "../../../editor/BlockNoteExtension.js";
 import { createBlockConfig, createBlockSpec } from "../../../schema/index.js";
 import {
   addDefaultPropsExternalHTML,
@@ -123,7 +123,7 @@ export const createCheckListItemBlockSpec = createBlockSpec(
     runsBefore: ["bulletListItem"],
   },
   [
-    createBlockNoteExtension({
+    createExtension({
       key: "check-list-item-shortcuts",
       keyboardShortcuts: {
         Enter: ({ editor }) => {
diff --git a/packages/core/src/blocks/ListItem/NumberedListItem/block.ts b/packages/core/src/blocks/ListItem/NumberedListItem/block.ts
index 70f1bcaa90..677d02523a 100644
--- a/packages/core/src/blocks/ListItem/NumberedListItem/block.ts
+++ b/packages/core/src/blocks/ListItem/NumberedListItem/block.ts
@@ -1,5 +1,5 @@
 import { getBlockInfoFromSelection } from "../../../api/getBlockInfoFromPos.js";
-import { createBlockNoteExtension } from "../../../editor/BlockNoteExtension.js";
+import { createExtension } from "../../../editor/BlockNoteExtension.js";
 import { createBlockConfig, createBlockSpec } from "../../../schema/index.js";
 import {
   addDefaultPropsExternalHTML,
@@ -91,7 +91,7 @@ export const createNumberedListItemBlockSpec = createBlockSpec(
     },
   },
   [
-    createBlockNoteExtension({
+    createExtension({
       key: "numbered-list-item-shortcuts",
       inputRules: [
         {
@@ -135,7 +135,7 @@ export const createNumberedListItemBlockSpec = createBlockSpec(
           return true;
         },
       },
-      plugins: [NumberedListIndexingDecorationPlugin()],
+      prosemirrorPlugins: [NumberedListIndexingDecorationPlugin()],
     }),
   ],
 );
diff --git a/packages/core/src/blocks/ListItem/ToggleListItem/block.ts b/packages/core/src/blocks/ListItem/ToggleListItem/block.ts
index e70da1b4b9..02b69d4e5b 100644
--- a/packages/core/src/blocks/ListItem/ToggleListItem/block.ts
+++ b/packages/core/src/blocks/ListItem/ToggleListItem/block.ts
@@ -1,4 +1,4 @@
-import { createBlockNoteExtension } from "../../../editor/BlockNoteExtension.js";
+import { createExtension } from "../../../editor/BlockNoteExtension.js";
 import { createBlockConfig, createBlockSpec } from "../../../schema/index.js";
 import {
   addDefaultPropsExternalHTML,
@@ -50,7 +50,7 @@ export const createToggleListItemBlockSpec = createBlockSpec(
     },
   },
   [
-    createBlockNoteExtension({
+    createExtension({
       key: "toggle-list-item-shortcuts",
       keyboardShortcuts: {
         Enter: ({ editor }) => {
diff --git a/packages/core/src/blocks/PageBreak/getPageBreakSlashMenuItems.ts b/packages/core/src/blocks/PageBreak/getPageBreakSlashMenuItems.ts
index 02502c5c5c..8edc3c5d62 100644
--- a/packages/core/src/blocks/PageBreak/getPageBreakSlashMenuItems.ts
+++ b/packages/core/src/blocks/PageBreak/getPageBreakSlashMenuItems.ts
@@ -1,6 +1,6 @@
 import { BlockNoteEditor } from "../../editor/BlockNoteEditor.js";
 import { DefaultSuggestionItem } from "../../extensions/SuggestionMenu/DefaultSuggestionItem.js";
-import { insertOrUpdateBlock } from "../../extensions/SuggestionMenu/getDefaultSlashMenuItems.js";
+import { insertOrUpdateBlockForSlashMenu } from "../../extensions/SuggestionMenu/getDefaultSlashMenuItems.js";
 import {
   BlockSchema,
   InlineContentSchema,
@@ -35,7 +35,7 @@ export function getPageBreakSlashMenuItems<
     items.push({
       ...editor.dictionary.slash_menu.page_break,
       onItemClick: () => {
-        insertOrUpdateBlock(editor, {
+        insertOrUpdateBlockForSlashMenu(editor, {
           type: "pageBreak",
         });
       },
diff --git a/packages/core/src/blocks/Paragraph/block.ts b/packages/core/src/blocks/Paragraph/block.ts
index 09a9cc9ddb..be52b6c9fc 100644
--- a/packages/core/src/blocks/Paragraph/block.ts
+++ b/packages/core/src/blocks/Paragraph/block.ts
@@ -1,4 +1,4 @@
-import { createBlockNoteExtension } from "../../editor/BlockNoteExtension.js";
+import { createExtension } from "../../editor/BlockNoteExtension.js";
 import { createBlockConfig, createBlockSpec } from "../../schema/index.js";
 import {
   addDefaultPropsExternalHTML,
@@ -55,7 +55,7 @@ export const createParagraphBlockSpec = createBlockSpec(
     runsBefore: ["default"],
   },
   [
-    createBlockNoteExtension({
+    createExtension({
       key: "paragraph-shortcuts",
       keyboardShortcuts: {
         "Mod-Alt-0": ({ editor }) => {
diff --git a/packages/core/src/blocks/Quote/block.ts b/packages/core/src/blocks/Quote/block.ts
index a0f6d6cb4a..bb14afbd77 100644
--- a/packages/core/src/blocks/Quote/block.ts
+++ b/packages/core/src/blocks/Quote/block.ts
@@ -1,4 +1,4 @@
-import { createBlockNoteExtension } from "../../editor/BlockNoteExtension.js";
+import { createExtension } from "../../editor/BlockNoteExtension.js";
 import { createBlockConfig, createBlockSpec } from "../../schema/index.js";
 import {
   addDefaultPropsExternalHTML,
@@ -54,7 +54,7 @@ export const createQuoteBlockSpec = createBlockSpec(
     },
   },
   [
-    createBlockNoteExtension({
+    createExtension({
       key: "quote-block-shortcuts",
       keyboardShortcuts: {
         "Mod-Alt-q": ({ editor }) => {
diff --git a/packages/core/src/blocks/Table/block.ts b/packages/core/src/blocks/Table/block.ts
index d101c5144f..18762d7cfe 100644
--- a/packages/core/src/blocks/Table/block.ts
+++ b/packages/core/src/blocks/Table/block.ts
@@ -2,7 +2,7 @@ import { Node, mergeAttributes } from "@tiptap/core";
 import { DOMParser, Fragment, Node as PMNode, Schema } from "prosemirror-model";
 import { CellSelection, TableView } from "prosemirror-tables";
 import { NodeView } from "prosemirror-view";
-import { createBlockNoteExtension } from "../../editor/BlockNoteExtension.js";
+import { createExtension } from "../../editor/BlockNoteExtension.js";
 import {
   BlockConfig,
   createBlockSpecFromTiptapNode,
@@ -385,7 +385,7 @@ export const createTableBlockSpec = () =>
     { node: TiptapTableNode, type: "table", content: "table" },
     tablePropSchema,
     [
-      createBlockNoteExtension({
+      createExtension({
         key: "table-extensions",
         tiptapExtensions: [
           TableExtension,
@@ -399,7 +399,7 @@ export const createTableBlockSpec = () =>
       // and all cells are selected. Uses a separate extension as it needs
       // priority over keyboard handlers in the `TableExtension`'s
       // `tableEditing` plugin.
-      createBlockNoteExtension({
+      createExtension({
         key: "table-keyboard-delete",
         keyboardShortcuts: {
           Backspace: ({ editor }) => {
diff --git a/packages/core/src/editor/BlockNoteEditor.test.ts b/packages/core/src/editor/BlockNoteEditor.test.ts
index 90c7cbcf4c..0e169f0aac 100644
--- a/packages/core/src/editor/BlockNoteEditor.test.ts
+++ b/packages/core/src/editor/BlockNoteEditor.test.ts
@@ -1,11 +1,12 @@
 import { expect, it } from "vitest";
+import * as Y from "yjs";
+
 import {
   getBlockInfo,
   getNearestBlockPos,
 } from "../api/getBlockInfoFromPos.js";
 import { BlockNoteEditor } from "./BlockNoteEditor.js";
-import { BlockNoteExtension } from "./BlockNoteExtension.js";
-import * as Y from "yjs";
+import { createExtension } from "./BlockNoteExtension.js";
 
 /**
  * @vitest-environment jsdom
@@ -131,18 +132,12 @@ it("onCreate event", () => {
   let created = false;
   BlockNoteEditor.create({
     extensions: [
-      (e) =>
-        new (class extends BlockNoteExtension {
-          public static key() {
-            return "test";
-          }
-          constructor(editor: BlockNoteEditor) {
-            super(editor);
-            editor.onCreate(() => {
-              created = true;
-            });
-          }
-        })(e),
+      createExtension({
+        key: "test",
+        mount: () => {
+          created = true;
+        },
+      }),
     ],
   });
   expect(created).toBe(true);
diff --git a/packages/core/src/editor/BlockNoteEditor.ts b/packages/core/src/editor/BlockNoteEditor.ts
index ed08960849..978964b3f0 100644
--- a/packages/core/src/editor/BlockNoteEditor.ts
+++ b/packages/core/src/editor/BlockNoteEditor.ts
@@ -1,22 +1,16 @@
 import {
-  AnyExtension,
   createDocument,
   EditorOptions,
-  Extension,
   FocusPosition,
   getSchema,
-  InputRule,
-  Mark,
   Editor as TiptapEditor,
-  Node as TipTapNode,
 } from "@tiptap/core";
 import { type Command, type Plugin, type Transaction } from "@tiptap/pm/state";
-import { dropCursor } from "prosemirror-dropcursor";
 import { Node, Schema } from "prosemirror-model";
 import * as Y from "yjs";
 
 import type { BlocksChanged } from "../api/getBlocksChangedByTransaction.js";
-import { editorHasBlockWithType } from "../blocks/defaultBlockTypeGuards.js";
+import { blockToNode } from "../api/nodeConversions/blockToNode.js";
 import {
   Block,
   BlockNoteSchema,
@@ -26,15 +20,7 @@ import {
   PartialBlock,
 } from "../blocks/index.js";
 import type { ThreadStore, User } from "../comments/index.js";
-import type { CommentsPlugin } from "../extensions/Comments/CommentsPlugin.js";
-import type { FilePanelProsemirrorPlugin } from "../extensions/FilePanel/FilePanelPlugin.js";
-import type { FormattingToolbarProsemirrorPlugin } from "../extensions/FormattingToolbar/FormattingToolbarPlugin.js";
-import type { LinkToolbarProsemirrorPlugin } from "../extensions/LinkToolbar/LinkToolbarPlugin.js";
-import type { ShowSelectionPlugin } from "../extensions/ShowSelection/ShowSelectionPlugin.js";
-import type { SideMenuProsemirrorPlugin } from "../extensions/SideMenu/SideMenuPlugin.js";
-import type { SuggestionMenuProseMirrorPlugin } from "../extensions/SuggestionMenu/SuggestionPlugin.js";
-import type { TableHandlesProsemirrorPlugin } from "../extensions/TableHandles/TableHandlesPlugin.js";
-import { UniqueID } from "../extensions/UniqueID/UniqueID.js";
+import { UniqueID } from "../extensions/tiptap-extensions/UniqueID/UniqueID.js";
 import type { Dictionary } from "../i18n/dictionary.js";
 import { en } from "../i18n/locales/index.js";
 import type {
@@ -50,16 +36,15 @@ import type {
   StyleSchema,
   StyleSpecs,
 } from "../schema/index.js";
+import "../style.css";
 import { mergeCSSClasses } from "../util/browser.js";
 import { EventEmitter } from "../util/EventEmitter.js";
 import type { NoInfer } from "../util/typescript.js";
-import { BlockNoteExtension } from "./BlockNoteExtension.js";
-import { getBlockNoteExtensions } from "./BlockNoteExtensions.js";
+import { Extension, ExtensionFactory } from "./BlockNoteExtension.js";
 import type { TextCursorPosition } from "./cursorPositionTypes.js";
 import {
   BlockManager,
   CollaborationManager,
-  type CollaborationOptions,
   EventManager,
   ExportManager,
   ExtensionManager,
@@ -70,36 +55,17 @@ import {
 import type { Selection } from "./selectionTypes.js";
 import { transformPasted } from "./transformPasted.js";
 
-import { updateBlockTr } from "../api/blockManipulation/commands/updateBlock/updateBlock.js";
-import { getBlockInfoFromTransaction } from "../api/getBlockInfoFromPos.js";
-import { blockToNode } from "../api/nodeConversions/blockToNode.js";
-import "../style.css";
-
-/**
- * A factory function that returns a BlockNoteExtension
- * This is useful so we can create extensions that require an editor instance
- * in the constructor
- */
-export type BlockNoteExtensionFactory = (
-  editor: BlockNoteEditor<any, any, any>,
-) => BlockNoteExtension;
-
-/**
- * We support Tiptap extensions and BlockNoteExtension based extensions
- */
-export type SupportedExtension = AnyExtension | BlockNoteExtension;
-
 export type BlockCache<
   BSchema extends BlockSchema = any,
   ISchema extends InlineContentSchema = any,
   SSchema extends StyleSchema = any,
 > = WeakMap<Node, Block<BSchema, ISchema, SSchema>>;
 
-export type BlockNoteEditorOptions<
+export interface BlockNoteEditorOptions<
   BSchema extends BlockSchema,
   ISchema extends InlineContentSchema,
   SSchema extends StyleSchema,
-> = {
+> {
   /**
    * Whether changes to blocks (like indentation, creating lists, changing headings) should be animated or not. Defaults to `true`.
    *
@@ -373,30 +339,15 @@ export type BlockNoteEditorOptions<
    */
   _tiptapOptions?: Partial<EditorOptions>;
 
-  /**
-   * (experimental) add extra extensions to the editor
-   *
-   * @deprecated, should use `extensions` instead
-   * @internal
-   */
-  _extensions?: Record<
-    string,
-    | { plugin: Plugin; priority?: number }
-    | ((editor: BlockNoteEditor<any, any, any>) => {
-        plugin: Plugin;
-        priority?: number;
-      })
-  >;
-
   /**
    * Register extensions to the editor.
    *
    * See [Extensions](/docs/features/extensions) for more info.
    *
-   * @remarks `BlockNoteExtension[]`
+   * @remarks `Extension[]`
    */
-  extensions?: Array<BlockNoteExtension | BlockNoteExtensionFactory>;
-};
+  extensions?: Array<Extension | ExtensionFactory>;
+}
 
 const blockNoteTipTapOptions = {
   enableInputRules: true,
@@ -417,9 +368,9 @@ export class BlockNoteEditor<
   public readonly pmSchema: Schema;
 
   /**
-   * extensions that are added to the editor, can be tiptap extensions or prosemirror plugins
+   * BlockNote extensions that are added to the editor, keyed by the extension key
    */
-  public extensions: Record<string, SupportedExtension> = {};
+  public extensions = new Map<string, Extension>();
 
   public readonly _tiptapEditor: TiptapEditor & {
     contentComponent: any;
@@ -453,56 +404,6 @@ export class BlockNoteEditor<
   public readonly inlineContentImplementations: InlineContentSpecs;
   public readonly styleImplementations: StyleSpecs;
 
-  public get formattingToolbar(): FormattingToolbarProsemirrorPlugin {
-    return this._extensionManager.formattingToolbar;
-  }
-
-  public get linkToolbar(): LinkToolbarProsemirrorPlugin<
-    BSchema,
-    ISchema,
-    SSchema
-  > {
-    return this._extensionManager.linkToolbar;
-  }
-
-  public get sideMenu(): SideMenuProsemirrorPlugin<BSchema, ISchema, SSchema> {
-    return this._extensionManager.sideMenu;
-  }
-
-  public get suggestionMenus(): SuggestionMenuProseMirrorPlugin<
-    BSchema,
-    ISchema,
-    SSchema
-  > {
-    return this._extensionManager.suggestionMenus;
-  }
-
-  public get filePanel():
-    | FilePanelProsemirrorPlugin<ISchema, SSchema>
-    | undefined {
-    return this._extensionManager.filePanel;
-  }
-
-  public get tableHandles():
-    | TableHandlesProsemirrorPlugin<ISchema, SSchema>
-    | undefined {
-    return this._extensionManager.tableHandles;
-  }
-
-  public get comments(): CommentsPlugin | undefined {
-    return this._collaborationManager?.comments;
-  }
-
-  public get showSelectionPlugin(): ShowSelectionPlugin {
-    return this._extensionManager.showSelectionPlugin;
-  }
-
-  /**
-   * The plugin for forking a document, only defined if in collaboration mode
-   */
-  public get forkYDocPlugin() {
-    return this._collaborationManager?.forkYDocPlugin;
-  }
   /**
    * The `uploadFile` method is what the editor uses when files need to be uploaded (for example when selecting an image to upload).
    * This method should set when creating the editor as this is application-specific.
@@ -554,30 +455,6 @@ export class BlockNoteEditor<
     >,
   ) {
     super();
-    const anyOpts = options as any;
-    if (anyOpts.onEditorContentChange) {
-      throw new Error(
-        "onEditorContentChange initialization option is deprecated, use <BlockNoteView onChange={...} />, the useEditorChange(...) hook, or editor.onChange(...)",
-      );
-    }
-
-    if (anyOpts.onTextCursorPositionChange) {
-      throw new Error(
-        "onTextCursorPositionChange initialization option is deprecated, use <BlockNoteView onSelectionChange={...} />, the useEditorSelectionChange(...) hook, or editor.onSelectionChange(...)",
-      );
-    }
-
-    if (anyOpts.onEditorReady) {
-      throw new Error(
-        "onEditorReady is deprecated. Editor is immediately ready for use after creation.",
-      );
-    }
-
-    if (anyOpts.editable) {
-      throw new Error(
-        "editable initialization option is deprecated, use <BlockNoteView editable={true/false} />, or alternatively editor.isEditable = true/false",
-      );
-    }
 
     this.dictionary = options.dictionary || en;
     this.settings = {
@@ -608,22 +485,22 @@ export class BlockNoteEditor<
 
     // Initialize CollaborationManager if collaboration is enabled or if comments are configured
     if (newOptions.collaboration || newOptions.comments) {
-      const collaborationOptions: CollaborationOptions = {
-        // Use collaboration options if available, otherwise provide defaults
-        fragment: newOptions.collaboration?.fragment || new Y.XmlFragment(),
-        user: newOptions.collaboration?.user || {
-          name: "User",
-          color: "#FF0000",
-        },
-        provider: newOptions.collaboration?.provider || null,
-        renderCursor: newOptions.collaboration?.renderCursor,
-        showCursorLabels: newOptions.collaboration?.showCursorLabels,
-        comments: newOptions.comments,
-        resolveUsers: newOptions.resolveUsers,
-      };
+      // const collaborationOptions: CollaborationOptions = {
+      //   // Use collaboration options if available, otherwise provide defaults
+      //   fragment: newOptions.collaboration?.fragment || new Y.XmlFragment(),
+      //   user: newOptions.collaboration?.user || {
+      //     name: "User",
+      //     color: "#FF0000",
+      //   },
+      //   provider: newOptions.collaboration?.provider || null,
+      //   renderCursor: newOptions.collaboration?.renderCursor,
+      //   showCursorLabels: newOptions.collaboration?.showCursorLabels,
+      //   comments: newOptions.comments,
+      //   resolveUsers: newOptions.resolveUsers,
+      // };
       this._collaborationManager = new CollaborationManager(
         this as any,
-        collaborationOptions,
+        newOptions,
       );
     } else {
       this._collaborationManager = undefined;
@@ -639,78 +516,7 @@ export class BlockNoteEditor<
     this.inlineContentImplementations = newOptions.schema.inlineContentSpecs;
     this.styleImplementations = newOptions.schema.styleSpecs;
 
-    this.extensions = {
-      ...getBlockNoteExtensions({
-        editor: this,
-        domAttributes: newOptions.domAttributes || {},
-        blockSpecs: this.schema.blockSpecs,
-        styleSpecs: this.schema.styleSpecs,
-        inlineContentSpecs: this.schema.inlineContentSpecs,
-        collaboration: newOptions.collaboration,
-        trailingBlock: newOptions.trailingBlock,
-        disableExtensions: newOptions.disableExtensions,
-        setIdAttribute: newOptions.setIdAttribute,
-        animations: newOptions.animations ?? true,
-        tableHandles: editorHasBlockWithType(this, "table"),
-        dropCursor: this.options.dropCursor ?? dropCursor,
-        placeholders: newOptions.placeholders,
-        tabBehavior: newOptions.tabBehavior,
-        pasteHandler: newOptions.pasteHandler,
-      }),
-      ...this._collaborationManager?.initExtensions(),
-    } as any;
-
-    // add extensions from _tiptapOptions
-    (newOptions._tiptapOptions?.extensions || []).forEach((ext) => {
-      this.extensions[ext.name] = ext;
-    });
-
-    // add extensions from options
-    for (let ext of newOptions.extensions || []) {
-      if (typeof ext === "function") {
-        // factory
-        ext = ext(this);
-      }
-      const key = (ext as any).key ?? (ext.constructor as any).key();
-      if (!key) {
-        throw new Error(
-          `Extension ${ext.constructor.name} does not have a key method`,
-        );
-      }
-      if (this.extensions[key]) {
-        throw new Error(
-          `Extension ${ext.constructor.name} already exists with key ${key}`,
-        );
-      }
-      this.extensions[key] = ext;
-    }
-
-    // (when passed in via the deprecated `_extensions` option)
-    Object.entries(newOptions._extensions || {}).forEach(([key, ext]) => {
-      // eslint-disable-next-line @typescript-eslint/no-this-alias
-      const editor = this;
-
-      const instance = typeof ext === "function" ? ext(editor) : ext;
-      if (!("plugin" in instance)) {
-        // Assume it is an Extension/Mark/Node
-        this.extensions[key] = instance;
-        return;
-      }
-
-      this.extensions[key] = new (class extends BlockNoteExtension {
-        public static key() {
-          return key;
-        }
-        constructor() {
-          super();
-          this.addProsemirrorPlugin(instance.plugin);
-        }
-        public get priority() {
-          return instance.priority;
-        }
-      })();
-    });
-
+    // TODO this should just be an extension
     if (newOptions.uploadFile) {
       const uploadFile = newOptions.uploadFile;
       this.uploadFile = async (file, blockId) => {
@@ -730,8 +536,7 @@ export class BlockNoteEditor<
     this.resolveFileUrl = newOptions.resolveFileUrl;
 
     const collaborationEnabled =
-      "ySyncPlugin" in this.extensions ||
-      "liveblocksExtension" in this.extensions;
+      "ySync" in this.extensions || "liveblocksExtension" in this.extensions;
 
     if (collaborationEnabled && newOptions.initialContent) {
       // eslint-disable-next-line no-console
@@ -740,105 +545,13 @@ export class BlockNoteEditor<
       );
     }
 
-    const blockExtensions = Object.fromEntries(
-      Object.values(this.schema.blockSpecs)
-        .map((block) => (block as any).extensions as any)
-        .filter((ext) => ext !== undefined)
-        .flat()
-        .map((ext) => [ext.key ?? ext.constructor.key(), ext]),
-    );
-    const tiptapExtensions = [
-      ...Object.entries({ ...this.extensions, ...blockExtensions }).map(
-        ([key, ext]) => {
-          if (
-            ext instanceof Extension ||
-            ext instanceof TipTapNode ||
-            ext instanceof Mark
-          ) {
-            // tiptap extension
-            return ext;
-          }
-
-          if (ext instanceof BlockNoteExtension) {
-            if (
-              !ext.plugins.length &&
-              !ext.keyboardShortcuts &&
-              !ext.inputRules &&
-              !ext.tiptapExtensions
-            ) {
-              return undefined;
-            }
-            // "blocknote" extensions (prosemirror plugins)
-            return Extension.create({
-              name: key,
-              priority: ext.priority,
-              addProseMirrorPlugins: () => ext.plugins,
-              addExtensions: () => ext.tiptapExtensions || [],
-              // TODO maybe collect all input rules from all extensions into one plugin
-              // TODO consider using the prosemirror-inputrules package instead
-              addInputRules: ext.inputRules
-                ? () =>
-                    ext.inputRules!.map(
-                      (inputRule) =>
-                        new InputRule({
-                          find: inputRule.find,
-                          handler: ({ range, match, state }) => {
-                            const replaceWith = inputRule.replace({
-                              match,
-                              range,
-                              editor: this,
-                            });
-                            if (replaceWith) {
-                              const cursorPosition =
-                                this.getTextCursorPosition();
-
-                              if (
-                                this.schema.blockSchema[
-                                  cursorPosition.block.type
-                                ].content !== "inline"
-                              ) {
-                                return undefined;
-                              }
-
-                              const blockInfo = getBlockInfoFromTransaction(
-                                state.tr,
-                              );
-                              const tr = state.tr.deleteRange(
-                                range.from,
-                                range.to,
-                              );
-
-                              updateBlockTr(
-                                tr,
-                                blockInfo.bnBlock.beforePos,
-                                replaceWith,
-                              );
-                              return undefined;
-                            }
-                            return null;
-                          },
-                        }),
-                    )
-                : undefined,
-              addKeyboardShortcuts: ext.keyboardShortcuts
-                ? () => {
-                    return Object.fromEntries(
-                      Object.entries(ext.keyboardShortcuts!).map(
-                        ([key, value]) => [
-                          key,
-                          () => value({ editor: this as any }),
-                        ],
-                      ),
-                    );
-                  }
-                : undefined,
-            });
-          }
-
-          return undefined;
-        },
-      ),
-    ].filter((ext): ext is Extension => ext !== undefined);
+    this._eventManager = new EventManager(this as any);
+    this._extensionManager = new ExtensionManager(this, newOptions);
+
+    const tiptapExtensions = this._extensionManager.getTiptapExtensions();
+
+    this.extensions = this._extensionManager.getExtensions();
+
     const tiptapOptions: EditorOptions = {
       ...blockNoteTipTapOptions,
       ...newOptions._tiptapOptions,
@@ -941,19 +654,9 @@ export class BlockNoteEditor<
     // Initialize managers
     this._blockManager = new BlockManager(this as any);
 
-    this._eventManager = new EventManager(this as any);
     this._exportManager = new ExportManager(this as any);
-    this._extensionManager = new ExtensionManager(this as any);
     this._selectionManager = new SelectionManager(this as any);
-    this._stateManager = new StateManager(
-      this as any,
-      collaborationEnabled
-        ? {
-            undo: this._collaborationManager?.getUndoCommand(),
-            redo: this._collaborationManager?.getRedoCommand(),
-          }
-        : undefined,
-    );
+    this._stateManager = new StateManager(this as any);
     this._styleManager = new StyleManager(this as any);
 
     this.emit("create");
@@ -962,7 +665,7 @@ export class BlockNoteEditor<
   // Manager instances
   private readonly _blockManager: BlockManager<any, any, any>;
   private readonly _collaborationManager?: CollaborationManager;
-  private readonly _eventManager: EventManager<any>;
+  private readonly _eventManager: EventManager<any, any, any>;
   private readonly _exportManager: ExportManager<any, any, any>;
   private readonly _extensionManager: ExtensionManager;
   private readonly _selectionManager: SelectionManager<any, any, any>;
@@ -1032,21 +735,25 @@ export class BlockNoteEditor<
     return this._stateManager.transact(callback);
   }
 
-  // TO DISCUSS
   /**
-   * Shorthand to get a typed extension from the editor, by
-   * just passing in the extension class.
-   *
-   * @param ext - The extension class to get
-   * @param key - optional, the key of the extension in the extensions object (defaults to the extension name)
-   * @returns The extension instance
+   * Remove an extension(s) from the editor
    */
-  public extension<T extends BlockNoteExtension>(
-    ext: { new (...args: any[]): T } & typeof BlockNoteExtension,
-    key = ext.key(),
-  ): T {
-    return this._extensionManager.extension(ext, key);
-  }
+  public removeExtension: ExtensionManager["unregisterExtension"] = (
+    ...args: Parameters<ExtensionManager["unregisterExtension"]>
+  ) => this._extensionManager.unregisterExtension(...args);
+
+  /**
+   * Register an extension to the editor
+   */
+  public registerExtension: ExtensionManager["registerExtension"] = (
+    ...args: Parameters<ExtensionManager["registerExtension"]>
+  ) => this._extensionManager.registerExtension(...args) as any;
+
+  /**
+   * Get an extension from the editor
+   */
+  public getExtension: ExtensionManager["getExtension"] = (...args) =>
+    this._extensionManager.getExtension(...args);
 
   /**
    * Mount the editor to a DOM element.
@@ -1099,6 +806,9 @@ export class BlockNoteEditor<
     return !this._tiptapEditor.isInitialized;
   }
 
+  /**
+   * Focus on the editor
+   */
   public focus() {
     if (this.headless) {
       return;
@@ -1106,6 +816,9 @@ export class BlockNoteEditor<
     this.prosemirrorView.focus();
   }
 
+  /**
+   * Blur the editor
+   */
   public blur() {
     if (this.headless) {
       return;
@@ -1113,6 +826,7 @@ export class BlockNoteEditor<
     this.prosemirrorView.dom.blur();
   }
 
+  // TODO move to extension
   public onUploadStart(callback: (blockId?: string) => void) {
     this.onUploadStartCallbacks.push(callback);
 
@@ -1367,14 +1081,14 @@ export class BlockNoteEditor<
   /**
    * Undo the last action.
    */
-  public undo() {
+  public undo(): boolean {
     return this._stateManager.undo();
   }
 
   /**
    * Redo the last action.
    */
-  public redo() {
+  public redo(): boolean {
     return this._stateManager.redo();
   }
 
@@ -1517,6 +1231,7 @@ export class BlockNoteEditor<
   ): string {
     return this._exportManager.blocksToFullHTML(blocks);
   }
+
   /**
    * Parses blocks from an HTML string. Tries to create `Block` objects out of any HTML block-level elements, and
    * `InlineNode` objects from any HTML inline elements, though not all element types are recognized. If BlockNote
@@ -1584,8 +1299,13 @@ export class BlockNoteEditor<
         getChanges(): BlocksChanged<BSchema, ISchema, SSchema>;
       },
     ) => void,
+    /**
+     * If true, the callback will be triggered when the changes are caused by a remote user
+     * @default true
+     */
+    includeUpdatesFromRemote?: boolean,
   ) {
-    return this._eventManager.onChange(callback);
+    return this._eventManager.onChange(callback, includeUpdatesFromRemote);
   }
 
   /**
@@ -1604,6 +1324,22 @@ export class BlockNoteEditor<
     );
   }
 
+  /**
+   * Gets the bounding box of a block.
+   * @param blockId The id of the block to get the bounding box of.
+   * @returns The bounding box of the block or undefined if the block is not found.
+   */
+  public getBlockClientRect(blockId: string): DOMRect | undefined {
+    const blockElement = this.prosemirrorView.root.querySelector(
+      `[data-node-type="blockContainer"][data-id="${blockId}"]`,
+    );
+    if (!blockElement) {
+      return;
+    }
+
+    return blockElement.getBoundingClientRect();
+  }
+
   /**
    * A callback function that runs when the editor has been initialized.
    *
@@ -1672,42 +1408,30 @@ export class BlockNoteEditor<
     );
   }
 
-  public openSuggestionMenu(
-    triggerCharacter: string,
-    pluginState?: {
-      deleteTriggerCharacter?: boolean;
-      ignoreQueryLength?: boolean;
-    },
-  ) {
-    if (!this.prosemirrorView || this.headless) {
-      return;
-    }
-
-    this.focus();
-    this.transact((tr) => {
-      if (pluginState?.deleteTriggerCharacter) {
-        tr.insertText(triggerCharacter);
-      }
-      tr.scrollIntoView().setMeta(this.suggestionMenus.plugins[0], {
-        triggerCharacter: triggerCharacter,
-        deleteTriggerCharacter: pluginState?.deleteTriggerCharacter || false,
-        ignoreQueryLength: pluginState?.ignoreQueryLength || false,
-      });
-    });
-  }
-
-  // `forceSelectionVisible` determines whether the editor selection is shows
-  // even when the editor is not focused. This is useful for e.g. creating new
-  // links, so the user still sees the affected content when an input field is
-  // focused.
-  // TODO: Reconsider naming?
-  public getForceSelectionVisible() {
-    return this.showSelectionPlugin.getEnabled();
-  }
-
-  public setForceSelectionVisible(forceSelectionVisible: boolean) {
-    this.showSelectionPlugin.setEnabled(forceSelectionVisible);
-  }
+  // TODO move to extension
+  // public openSuggestionMenu(
+  //   triggerCharacter: string,
+  //   pluginState?: {
+  //     deleteTriggerCharacter?: boolean;
+  //     ignoreQueryLength?: boolean;
+  //   },
+  // ) {
+  //   if (!this.prosemirrorView) {
+  //     return;
+  //   }
+
+  //   this.focus();
+  //   this.transact((tr) => {
+  //     if (pluginState?.deleteTriggerCharacter) {
+  //       tr.insertText(triggerCharacter);
+  //     }
+  //     tr.scrollIntoView().setMeta(this.suggestionMenus.plugins[0], {
+  //       triggerCharacter: triggerCharacter,
+  //       deleteTriggerCharacter: pluginState?.deleteTriggerCharacter || false,
+  //       ignoreQueryLength: pluginState?.ignoreQueryLength || false,
+  //     });
+  //   });
+  // }
 
   /**
    * Paste HTML into the editor. Defaults to converting HTML to BlockNote HTML.
diff --git a/packages/core/src/editor/BlockNoteExtension.ts b/packages/core/src/editor/BlockNoteExtension.ts
index f56d6c736a..fe1e78a952 100644
--- a/packages/core/src/editor/BlockNoteExtension.ts
+++ b/packages/core/src/editor/BlockNoteExtension.ts
@@ -1,42 +1,59 @@
-import { Plugin } from "prosemirror-state";
-import { EventEmitter } from "../util/EventEmitter.js";
-
-import { AnyExtension } from "@tiptap/core";
-import {
-  BlockSchema,
-  InlineContentSchema,
-  PartialBlockNoDefaults,
-  StyleSchema,
-} from "../schema/index.js";
-import { BlockNoteEditor } from "./BlockNoteEditor.js";
-
-export abstract class BlockNoteExtension<
-  TEvent extends Record<string, any> = any,
-> extends EventEmitter<TEvent> {
-  public static key(): string {
-    throw new Error("You must implement the key method in your extension");
-  }
+import { Store, StoreOptions } from "@tanstack/store";
+import { type AnyExtension } from "@tiptap/core";
+import type { Plugin as ProsemirrorPlugin } from "prosemirror-state";
+import type { PartialBlockNoDefaults } from "../schema/index.js";
+import type {
+  BlockNoteEditor,
+  BlockNoteEditorOptions,
+} from "./BlockNoteEditor.js";
 
-  protected addProsemirrorPlugin(plugin: Plugin) {
-    this.plugins.push(plugin);
-  }
+/**
+ * This function is called when the extension is destroyed.
+ */
+type OnDestroy = () => void;
 
-  public readonly plugins: Plugin[] = [];
-  public get priority(): number | undefined {
-    return undefined;
-  }
+/**
+ * Describes a BlockNote extension.
+ */
+export interface Extension<State = any, Key extends string = string> {
+  /**
+   * The unique identifier for the extension.
+   */
+  readonly key: Key;
 
-  // eslint-disable-next-line
-  constructor(..._args: any[]) {
-    super();
-    // Allow subclasses to have constructors with parameters
-    // without this, we can't easily implement BlockNoteEditor.extension(MyExtension) pattern
-  }
+  /**
+   * Triggered when the extension is mounted to the editor.
+   */
+  readonly mount?: (ctx: {
+    /**
+     * The DOM element that the editor is mounted to.
+     */
+    dom: HTMLElement;
+    /**
+     * The root document of the {@link document} that the editor is mounted to.
+     */
+    root: Document | ShadowRoot;
+    /**
+     * An {@link AbortSignal} that will be aborted when the extension is destroyed.
+     */
+    signal: AbortSignal;
+  }) => void | OnDestroy;
 
   /**
-   * Input rules for the block
+   * The store for the extension.
    */
-  public inputRules?: InputRule[];
+  readonly store?: Store<State>;
+
+  /**
+   * Declares what {@link Extension}s that this extension depends on.
+   */
+  readonly runsBefore?: ReadonlyArray<string>;
+
+  /**
+   * Input rules for a block: An input rule is what is used to replace text in a block when a regular expression match is found.
+   * As an example, typing `#` in a paragraph block will trigger an input rule to replace the text with a heading block.
+   */
+  readonly inputRules?: ReadonlyArray<InputRule>;
 
   /**
    * A mapping of a keyboard shortcut to a function that will be called when the shortcut is pressed
@@ -60,17 +77,27 @@ export abstract class BlockNoteExtension<
    * }
    * ```
    */
-  public keyboardShortcuts?: Record<
+  readonly keyboardShortcuts?: Record<
     string,
-    (ctx: {
-      editor: BlockNoteEditor<BlockSchema, InlineContentSchema, StyleSchema>;
-    }) => boolean
+    (ctx: { editor: BlockNoteEditor<any, any, any> }) => boolean
   >;
 
-  public tiptapExtensions?: AnyExtension[];
+  /**
+   * Add additional prosemirror plugins to the editor.
+   */
+  readonly prosemirrorPlugins?: ReadonlyArray<ProsemirrorPlugin>;
+
+  /**
+   * Add additional tiptap extensions to the editor.
+   */
+  readonly tiptapExtensions?: ReadonlyArray<AnyExtension>;
 }
 
-export type InputRule = {
+/**
+ * An input rule is what is used to replace text in a block when a regular expression match is found.
+ * As an example, typing `#` in a paragraph block will trigger an input rule to replace the text with a heading block.
+ */
+type InputRule = {
   /**
    * The regex to match when to trigger the input rule
    */
@@ -96,23 +123,83 @@ export type InputRule = {
   }) => undefined | PartialBlockNoDefaults<any, any, any>;
 };
 
+export type ExtensionFactory<
+  State = any,
+  Options extends BlockNoteEditorOptions<
+    any,
+    any,
+    any
+  > = BlockNoteEditorOptions<any, any, any>,
+  Key extends string = string,
+> = (
+  editor: BlockNoteEditor<any, any, any>,
+  options: Options,
+) => Extension<State, Key> | undefined;
+
+export type ExtractExtensionKey<T extends ExtensionFactory | undefined> =
+  T extends (
+    editor: BlockNoteEditor<any, any, any>,
+    options: BlockNoteEditorOptions<any, any, any>,
+  ) => Extension<any, infer Key> | undefined
+    ? Key
+    : never;
+
+export type ExtractExtensionByKey<
+  T extends ExtensionFactory<any, any, any>,
+  Key extends string,
+> = T extends (
+  editor: BlockNoteEditor<any, any, any>,
+  options: BlockNoteEditorOptions<any, any, any>,
+) => Extension<any, Key> | undefined
+  ? ReturnType<T>
+  : never;
+
+// a type that maps the extension key to the return type of the extension factory
+export type ExtensionMap<
+  T extends ReadonlyArray<ExtensionFactory | undefined>,
+> = {
+  [K in ExtractExtensionKey<T[number]>]: ExtractExtensionByKey<
+    Exclude<T[number], undefined>,
+    K
+  >;
+};
+
 /**
- * This creates an instance of a BlockNoteExtension that can be used to add to a schema.
- * It is a bit of a hack, but it works.
+ * Helper function to create a BlockNote extension.
+ * Can accept either an ExtensionFactory (function) or an Extension (object).
+ * If an Extension is provided, it will be wrapped in a factory function.
  */
-export function createBlockNoteExtension(
-  options: Partial<
-    Pick<
-      BlockNoteExtension,
-      "inputRules" | "keyboardShortcuts" | "plugins" | "tiptapExtensions"
-    >
-  > & { key: string },
-) {
-  const x = Object.create(BlockNoteExtension.prototype);
-  x.key = options.key;
-  x.inputRules = options.inputRules;
-  x.keyboardShortcuts = options.keyboardShortcuts;
-  x.plugins = options.plugins ?? [];
-  x.tiptapExtensions = options.tiptapExtensions;
-  return x as BlockNoteExtension;
+export function createExtension<const T extends ExtensionFactory>(ext: T): T;
+export function createExtension<const Key extends string, const State = any>(
+  ext: Extension<State, Key>,
+): () => Extension<State, Key>;
+export function createExtension<
+  const T extends ExtensionFactory | Extension | undefined = any,
+>(
+  extension: T,
+): T extends Extension
+  ? () => T
+  : T extends ExtensionFactory
+    ? T
+    : () => undefined {
+  if (typeof extension === "function") {
+    return extension as any;
+  }
+
+  if (typeof extension === "object" && "key" in extension) {
+    return (() => extension) as any;
+  }
+
+  if (!extension) {
+    return (() => undefined) as any;
+  }
+
+  throw new Error("Invalid extension", { cause: { extension } });
+}
+
+export function createStore<T = any>(
+  initialState: T,
+  options?: StoreOptions<T>,
+): Store<T> {
+  return new Store(initialState, options);
 }
diff --git a/packages/core/src/editor/BlockNoteExtensions.ts b/packages/core/src/editor/BlockNoteExtensions.ts
deleted file mode 100644
index 754c90b428..0000000000
--- a/packages/core/src/editor/BlockNoteExtensions.ts
+++ /dev/null
@@ -1,325 +0,0 @@
-import { AnyExtension, Extension, extensions, Node } from "@tiptap/core";
-import { Gapcursor } from "@tiptap/extension-gapcursor";
-import { History } from "@tiptap/extension-history";
-import { Link } from "@tiptap/extension-link";
-import { Text } from "@tiptap/extension-text";
-import { Plugin } from "prosemirror-state";
-import * as Y from "yjs";
-
-import { createDropFileExtension } from "../api/clipboard/fromClipboard/fileDropExtension.js";
-import { createPasteFromClipboardExtension } from "../api/clipboard/fromClipboard/pasteExtension.js";
-import { createCopyToClipboardExtension } from "../api/clipboard/toClipboard/copyExtension.js";
-import type { ThreadStore, User } from "../comments/index.js";
-import { BackgroundColorExtension } from "../extensions/BackgroundColor/BackgroundColorExtension.js";
-import { BlockChangePlugin } from "../extensions/BlockChange/BlockChangePlugin.js";
-import { CursorPlugin } from "../extensions/Collaboration/CursorPlugin.js";
-import { ForkYDocPlugin } from "../extensions/Collaboration/ForkYDocPlugin.js";
-import { SyncPlugin } from "../extensions/Collaboration/SyncPlugin.js";
-import { UndoPlugin } from "../extensions/Collaboration/UndoPlugin.js";
-import { SchemaMigrationPlugin } from "../extensions/Collaboration/schemaMigration/SchemaMigrationPlugin.js";
-import { CommentMark } from "../extensions/Comments/CommentMark.js";
-import { CommentsPlugin } from "../extensions/Comments/CommentsPlugin.js";
-import { FilePanelProsemirrorPlugin } from "../extensions/FilePanel/FilePanelPlugin.js";
-import { FormattingToolbarProsemirrorPlugin } from "../extensions/FormattingToolbar/FormattingToolbarPlugin.js";
-import { HardBreak } from "../extensions/HardBreak/HardBreak.js";
-import { KeyboardShortcutsExtension } from "../extensions/KeyboardShortcuts/KeyboardShortcutsExtension.js";
-import { LinkToolbarProsemirrorPlugin } from "../extensions/LinkToolbar/LinkToolbarPlugin.js";
-import {
-  DEFAULT_LINK_PROTOCOL,
-  VALID_LINK_PROTOCOLS,
-} from "../extensions/LinkToolbar/protocols.js";
-import { NodeSelectionKeyboardPlugin } from "../extensions/NodeSelectionKeyboard/NodeSelectionKeyboardPlugin.js";
-import { PlaceholderPlugin } from "../extensions/Placeholder/PlaceholderPlugin.js";
-import { PreviousBlockTypePlugin } from "../extensions/PreviousBlockType/PreviousBlockTypePlugin.js";
-import { ShowSelectionPlugin } from "../extensions/ShowSelection/ShowSelectionPlugin.js";
-import { SideMenuProsemirrorPlugin } from "../extensions/SideMenu/SideMenuPlugin.js";
-import { SuggestionMenuProseMirrorPlugin } from "../extensions/SuggestionMenu/SuggestionPlugin.js";
-import {
-  SuggestionAddMark,
-  SuggestionDeleteMark,
-  SuggestionModificationMark,
-} from "../extensions/Suggestions/SuggestionMarks.js";
-import { TableHandlesProsemirrorPlugin } from "../extensions/TableHandles/TableHandlesPlugin.js";
-import { TextAlignmentExtension } from "../extensions/TextAlignment/TextAlignmentExtension.js";
-import { TextColorExtension } from "../extensions/TextColor/TextColorExtension.js";
-import { TrailingNode } from "../extensions/TrailingNode/TrailingNodeExtension.js";
-import UniqueID from "../extensions/UniqueID/UniqueID.js";
-import { BlockContainer, BlockGroup, Doc } from "../pm-nodes/index.js";
-import {
-  BlockNoteDOMAttributes,
-  BlockSchema,
-  BlockSpecs,
-  InlineContentSchema,
-  InlineContentSpecs,
-  StyleSchema,
-  StyleSpecs,
-} from "../schema/index.js";
-import type {
-  BlockNoteEditor,
-  BlockNoteEditorOptions,
-  SupportedExtension,
-} from "./BlockNoteEditor.js";
-import { BlockNoteSchema } from "../blocks/BlockNoteSchema.js";
-
-type ExtensionOptions<
-  BSchema extends BlockSchema,
-  I extends InlineContentSchema,
-  S extends StyleSchema,
-> = {
-  editor: BlockNoteEditor<BSchema, I, S>;
-  domAttributes: Partial<BlockNoteDOMAttributes>;
-  blockSpecs: BlockSpecs;
-  inlineContentSpecs: InlineContentSpecs;
-  styleSpecs: StyleSpecs;
-  trailingBlock: boolean | undefined;
-  collaboration?: {
-    fragment: Y.XmlFragment;
-    user: {
-      name: string;
-      color: string;
-      [key: string]: string;
-    };
-    provider: any;
-    renderCursor?: (user: any) => HTMLElement;
-    showCursorLabels?: "always" | "activity";
-  };
-  disableExtensions: string[] | undefined;
-  setIdAttribute?: boolean;
-  animations: boolean;
-  tableHandles: boolean;
-  dropCursor: (opts: any) => Plugin;
-  placeholders: Record<
-    string | "default" | "emptyDocument",
-    string | undefined
-  >;
-  tabBehavior?: "prefer-navigate-ui" | "prefer-indent";
-  comments?: {
-    schema?: BlockNoteSchema<any, any, any>;
-    threadStore: ThreadStore;
-    resolveUsers?: (userIds: string[]) => Promise<User[]>;
-  };
-  pasteHandler: BlockNoteEditorOptions<any, any, any>["pasteHandler"];
-};
-
-/**
- * Get all the Tiptap extensions BlockNote is configured with by default
- */
-export const getBlockNoteExtensions = <
-  BSchema extends BlockSchema,
-  I extends InlineContentSchema,
-  S extends StyleSchema,
->(
-  opts: ExtensionOptions<BSchema, I, S>,
-) => {
-  const ret: Record<string, SupportedExtension> = {};
-  const tiptapExtensions = getTipTapExtensions(opts);
-
-  for (const ext of tiptapExtensions) {
-    ret[ext.name] = ext;
-  }
-
-  if (opts.collaboration) {
-    ret["ySyncPlugin"] = new SyncPlugin(opts.collaboration.fragment);
-    ret["yUndoPlugin"] = new UndoPlugin({ editor: opts.editor });
-
-    if (opts.collaboration.provider?.awareness) {
-      ret["yCursorPlugin"] = new CursorPlugin(opts.collaboration);
-    }
-    ret["forkYDocPlugin"] = new ForkYDocPlugin({
-      editor: opts.editor,
-      collaboration: opts.collaboration,
-    });
-    ret["schemaMigrationPlugin"] = new SchemaMigrationPlugin(
-      opts.collaboration.fragment,
-    );
-  }
-
-  // Note: this is pretty hardcoded and will break when user provides plugins with same keys.
-  // Define name on plugins instead and not make this a map?
-  ret["formattingToolbar"] = new FormattingToolbarProsemirrorPlugin(
-    opts.editor,
-  );
-  ret["linkToolbar"] = new LinkToolbarProsemirrorPlugin(opts.editor);
-  ret["sideMenu"] = new SideMenuProsemirrorPlugin(opts.editor);
-  ret["suggestionMenus"] = new SuggestionMenuProseMirrorPlugin(opts.editor);
-  ret["filePanel"] = new FilePanelProsemirrorPlugin(opts.editor as any);
-  ret["placeholder"] = new PlaceholderPlugin(opts.editor, opts.placeholders);
-
-  if (opts.animations ?? true) {
-    ret["animations"] = new PreviousBlockTypePlugin();
-  }
-
-  if (opts.tableHandles) {
-    ret["tableHandles"] = new TableHandlesProsemirrorPlugin(opts.editor as any);
-  }
-
-  ret["nodeSelectionKeyboard"] = new NodeSelectionKeyboardPlugin();
-  ret["blockChange"] = new BlockChangePlugin();
-
-  ret["showSelection"] = new ShowSelectionPlugin(opts.editor);
-
-  if (opts.comments) {
-    ret["comments"] = new CommentsPlugin(
-      opts.editor,
-      opts.comments.threadStore,
-      CommentMark.name,
-      opts.comments.resolveUsers,
-      opts.comments.schema,
-    );
-  }
-
-  const disableExtensions: string[] = opts.disableExtensions || [];
-  for (const ext of disableExtensions) {
-    delete ret[ext];
-  }
-
-  return ret;
-};
-
-let LINKIFY_INITIALIZED = false;
-
-/**
- * Get all the Tiptap extensions BlockNote is configured with by default
- */
-const getTipTapExtensions = <
-  BSchema extends BlockSchema,
-  I extends InlineContentSchema,
-  S extends StyleSchema,
->(
-  opts: ExtensionOptions<BSchema, I, S>,
-) => {
-  const tiptapExtensions: AnyExtension[] = [
-    extensions.ClipboardTextSerializer,
-    extensions.Commands,
-    extensions.Editable,
-    extensions.FocusEvents,
-    extensions.Tabindex,
-
-    // DevTools,
-    Gapcursor,
-
-    // DropCursor,
-    Extension.create({
-      name: "dropCursor",
-      addProseMirrorPlugins: () => [
-        opts.dropCursor({
-          width: 5,
-          color: "#ddeeff",
-          editor: opts.editor,
-        }),
-      ],
-    }),
-
-    UniqueID.configure({
-      // everything from bnBlock group (nodes that represent a BlockNote block should have an id)
-      types: ["blockContainer", "columnList", "column"],
-      setIdAttribute: opts.setIdAttribute,
-    }),
-    HardBreak,
-    // Comments,
-
-    // basics:
-    Text,
-
-    // marks:
-    SuggestionAddMark,
-    SuggestionDeleteMark,
-    SuggestionModificationMark,
-    Link.extend({
-      inclusive: false,
-    }).configure({
-      defaultProtocol: DEFAULT_LINK_PROTOCOL,
-      // only call this once if we have multiple editors installed. Or fix https://github.com/ueberdosis/tiptap/issues/5450
-      protocols: LINKIFY_INITIALIZED ? [] : VALID_LINK_PROTOCOLS,
-    }),
-    ...(Object.values(opts.styleSpecs).map((styleSpec) => {
-      return styleSpec.implementation.mark.configure({
-        editor: opts.editor as any,
-      });
-    }) as any[]),
-
-    TextColorExtension,
-
-    BackgroundColorExtension,
-    TextAlignmentExtension,
-
-    // make sure escape blurs editor, so that we can tab to other elements in the host page (accessibility)
-    Extension.create({
-      name: "OverrideEscape",
-      addKeyboardShortcuts() {
-        return {
-          Escape: () => {
-            if (opts.editor.suggestionMenus.shown) {
-              // escape is handled by suggestionmenu
-              return false;
-            }
-            return this.editor.commands.blur();
-          },
-        };
-      },
-    }),
-
-    // nodes
-    Doc,
-    BlockContainer.configure({
-      editor: opts.editor,
-      domAttributes: opts.domAttributes,
-    }),
-    KeyboardShortcutsExtension.configure({
-      editor: opts.editor,
-      tabBehavior: opts.tabBehavior,
-    }),
-    BlockGroup.configure({
-      domAttributes: opts.domAttributes,
-    }),
-    ...Object.values(opts.inlineContentSpecs)
-      .filter((a) => a.config !== "link" && a.config !== "text")
-      .map((inlineContentSpec) => {
-        return inlineContentSpec.implementation!.node.configure({
-          editor: opts.editor as any,
-        });
-      }),
-
-    ...Object.values(opts.blockSpecs).flatMap((blockSpec) => {
-      return [
-        // the node extension implementations
-        ...("node" in blockSpec.implementation
-          ? [
-              (blockSpec.implementation.node as Node).configure({
-                editor: opts.editor,
-                domAttributes: opts.domAttributes,
-              }),
-            ]
-          : []),
-      ];
-    }),
-    createCopyToClipboardExtension(opts.editor),
-    createPasteFromClipboardExtension(
-      opts.editor,
-      opts.pasteHandler ||
-        ((context: {
-          defaultPasteHandler: (context?: {
-            prioritizeMarkdownOverHTML?: boolean;
-            plainTextAsMarkdown?: boolean;
-          }) => boolean | undefined;
-        }) => context.defaultPasteHandler()),
-    ),
-    createDropFileExtension(opts.editor),
-
-    // This needs to be at the bottom of this list, because Key events (such as enter, when selecting a /command),
-    // should be handled before Enter handlers in other components like splitListItem
-    ...(opts.trailingBlock === undefined || opts.trailingBlock
-      ? [TrailingNode]
-      : []),
-    ...(opts.comments ? [CommentMark] : []),
-  ];
-
-  LINKIFY_INITIALIZED = true;
-
-  if (!opts.collaboration) {
-    // disable history extension when collaboration is enabled as y-prosemirror takes care of undo / redo
-    tiptapExtensions.push(History);
-  }
-
-  return tiptapExtensions;
-};
diff --git a/packages/core/src/editor/managers/CollaborationManager.ts b/packages/core/src/editor/managers/CollaborationManager.ts
index 8273fb5cb4..7d25cfe6f5 100644
--- a/packages/core/src/editor/managers/CollaborationManager.ts
+++ b/packages/core/src/editor/managers/CollaborationManager.ts
@@ -1,55 +1,11 @@
-import * as Y from "yjs";
-import { redoCommand, undoCommand } from "y-prosemirror";
-import { CommentsPlugin } from "../../extensions/Comments/CommentsPlugin.js";
-import { CommentMark } from "../../extensions/Comments/CommentMark.js";
-import { ForkYDocPlugin } from "../../extensions/Collaboration/ForkYDocPlugin.js";
-import { SyncPlugin } from "../../extensions/Collaboration/SyncPlugin.js";
-import { UndoPlugin } from "../../extensions/Collaboration/UndoPlugin.js";
-import { CursorPlugin } from "../../extensions/Collaboration/CursorPlugin.js";
-import type { ThreadStore, User } from "../../comments/index.js";
-import type { BlockNoteEditor } from "../BlockNoteEditor.js";
-import { CustomBlockNoteSchema } from "../../schema/schema.js";
+import type { User } from "../../comments/index.js";
+import { YCursor } from "../../extensions/Collaboration/YCursorPlugin.js";
+import type {
+  BlockNoteEditor,
+  BlockNoteEditorOptions,
+} from "../BlockNoteEditor.js";
 
-export interface CollaborationOptions {
-  /**
-   * The Yjs XML fragment that's used for collaboration.
-   */
-  fragment: Y.XmlFragment;
-  /**
-   * The user info for the current user that's shown to other collaborators.
-   */
-  user: {
-    name: string;
-    color: string;
-  };
-  /**
-   * A Yjs provider (used for awareness / cursor information)
-   * Can be null for comments-only mode
-   */
-  provider: any;
-  /**
-   * Optional function to customize how cursors of users are rendered
-   */
-  renderCursor?: (user: any) => HTMLElement;
-  /**
-   * Optional flag to set when the user label should be shown with the default
-   * collaboration cursor. Setting to "always" will always show the label,
-   * while "activity" will only show the label when the user moves the cursor
-   * or types. Defaults to "activity".
-   */
-  showCursorLabels?: "always" | "activity";
-  /**
-   * Comments configuration - can be used with or without collaboration
-   */
-  comments?: {
-    schema?: CustomBlockNoteSchema<any, any, any>;
-    threadStore: ThreadStore;
-  };
-  /**
-   * Function to resolve user IDs to user objects - required for comments
-   */
-  resolveUsers?: (userIds: string[]) => Promise<User[]>;
-}
+// TODO remove this manager completely
 
 /**
  * CollaborationManager handles all collaboration-related functionality
@@ -57,98 +13,21 @@ export interface CollaborationOptions {
  */
 export class CollaborationManager {
   private editor: BlockNoteEditor;
-  private options: CollaborationOptions;
-  private _commentsPlugin?: CommentsPlugin;
-  private _forkYDocPlugin?: ForkYDocPlugin;
-  private _syncPlugin?: SyncPlugin;
-  private _undoPlugin?: UndoPlugin;
-  private _cursorPlugin?: CursorPlugin;
+  private options: BlockNoteEditorOptions<any, any, any>;
 
-  constructor(editor: BlockNoteEditor, options: CollaborationOptions) {
+  constructor(
+    editor: BlockNoteEditor,
+    options: BlockNoteEditorOptions<any, any, any>,
+  ) {
     this.editor = editor;
     this.options = options;
   }
 
-  /**
-   * Get the sync plugin instance
-   */
-  public get syncPlugin(): SyncPlugin | undefined {
-    return this._syncPlugin;
-  }
-
-  /**
-   * Get the undo plugin instance
-   */
-  public get undoPlugin(): UndoPlugin | undefined {
-    return this._undoPlugin;
-  }
-
-  /**
-   * Get the cursor plugin instance
-   */
-  public get cursorPlugin(): CursorPlugin | undefined {
-    return this._cursorPlugin;
-  }
-
-  /**
-   * Get the fork YDoc plugin instance
-   */
-  public get forkYDocPlugin(): ForkYDocPlugin | undefined {
-    return this._forkYDocPlugin;
-  }
-
-  // Initialize collaboration plugins
-  public initExtensions(): Record<string, unknown> {
-    // Only create collaboration plugins when real-time collaboration is enabled
-    const extensions: Record<string, unknown> = {};
-
-    // Initialize sync plugin
-    this._syncPlugin = new SyncPlugin(this.options.fragment);
-    extensions.ySyncPlugin = this._syncPlugin;
-
-    // Initialize undo plugin
-    this._undoPlugin = new UndoPlugin({ editor: this.editor });
-    extensions.yUndoPlugin = this._undoPlugin;
-
-    // Initialize cursor plugin if provider has awareness
-    if (this.options.provider?.awareness) {
-      this._cursorPlugin = new CursorPlugin(this.options);
-      extensions.yCursorPlugin = this._cursorPlugin;
-    }
-
-    // Initialize fork YDoc plugin
-    this._forkYDocPlugin = new ForkYDocPlugin({
-      editor: this.editor,
-      collaboration: this.options,
-    });
-    extensions.forkYDocPlugin = this._forkYDocPlugin;
-
-    if (this.options.comments) {
-      if (!this.options.resolveUsers) {
-        throw new Error("resolveUsers is required when using comments");
-      }
-
-      // Create CommentsPlugin instance and add it to editor extensions
-      this._commentsPlugin = new CommentsPlugin(
-        this.editor,
-        this.options.comments.threadStore,
-        CommentMark.name,
-        this.options.resolveUsers,
-        this.options.comments.schema,
-      );
-
-      // Add the comments plugin to the editor's extensions
-      extensions.comments = this._commentsPlugin;
-      extensions.commentMark = CommentMark;
-    }
-    return extensions;
-  }
-
   /**
    * Update the user info for the current user that's shown to other collaborators
    */
   public updateUserInfo(user: { name: string; color: string }) {
-    const cursor = this.cursorPlugin;
+    const cursor = this.editor.getExtension(YCursor);
     if (!cursor) {
       throw new Error(
         "Cannot update collaboration user info when collaboration is disabled.",
@@ -157,41 +36,13 @@ export class CollaborationManager {
     cursor.updateUser(user);
   }
 
-  /**
-   * Get the collaboration undo command
-   */
-  public getUndoCommand() {
-    return undoCommand;
-  }
-
-  /**
-   * Get the collaboration redo command
-   */
-  public getRedoCommand() {
-    return redoCommand;
-  }
-
   /**
    * Check if initial content should be avoided due to collaboration
    */
   public shouldAvoidInitialContent(): boolean {
     // Only avoid initial content when real-time collaboration is enabled
-    // (i.e., when we have a provider)
-    return !!this.options.provider;
-  }
-
-  /**
-   * Get the collaboration options
-   */
-  public getOptions(): CollaborationOptions {
-    return this.options;
-  }
-
-  /**
-   * Get the comments plugin if available
-   */
-  public get comments(): CommentsPlugin | undefined {
-    return this._commentsPlugin;
+    // (i.e., when we have a fragment)
+    return !!this.options.collaboration?.fragment;
   }
 
   /**
diff --git a/packages/core/src/editor/managers/EventManager.ts b/packages/core/src/editor/managers/EventManager.ts
index abb2d1f3fb..65a1349990 100644
--- a/packages/core/src/editor/managers/EventManager.ts
+++ b/packages/core/src/editor/managers/EventManager.ts
@@ -5,6 +5,11 @@ import {
 } from "../../api/getBlocksChangedByTransaction.js";
 import { Transaction } from "prosemirror-state";
 import { EventEmitter } from "../../util/EventEmitter.js";
+import {
+  BlockSchema,
+  InlineContentSchema,
+  StyleSchema,
+} from "../../schema/index.js";
 
 /**
  * A function that can be used to unsubscribe from an event.
@@ -14,22 +19,25 @@ export type Unsubscribe = () => void;
 /**
  * EventManager is a class which manages the events of the editor
  */
-export class EventManager<Editor extends BlockNoteEditor> extends EventEmitter<{
+export class EventManager<
+  BSchema extends BlockSchema,
+  I extends InlineContentSchema,
+  S extends StyleSchema,
+> extends EventEmitter<{
   onChange: [
-    editor: Editor,
     ctx: {
-      getChanges(): BlocksChanged<
-        Editor["schema"]["blockSchema"],
-        Editor["schema"]["inlineContentSchema"],
-        Editor["schema"]["styleSchema"]
-      >;
+      editor: BlockNoteEditor<BSchema, I, S>;
+      transaction: Transaction;
+      appendedTransactions: Transaction[];
     },
   ];
-  onSelectionChange: [ctx: { editor: Editor; transaction: Transaction }];
-  onMount: [ctx: { editor: Editor }];
-  onUnmount: [ctx: { editor: Editor }];
+  onSelectionChange: [
+    ctx: { editor: BlockNoteEditor<BSchema, I, S>; transaction: Transaction },
+  ];
+  onMount: [ctx: { editor: BlockNoteEditor<BSchema, I, S> }];
+  onUnmount: [ctx: { editor: BlockNoteEditor<BSchema, I, S> }];
 }> {
-  constructor(private editor: Editor) {
+  constructor(private editor: BlockNoteEditor<BSchema, I, S>) {
     super();
     // We register tiptap events only once the editor is finished initializing
     // otherwise we would be trying to register events on a tiptap editor which does not exist yet
@@ -37,14 +45,7 @@ export class EventManager<Editor extends BlockNoteEditor> extends EventEmitter<{
       editor._tiptapEditor.on(
         "update",
         ({ transaction, appendedTransactions }) => {
-          this.emit("onChange", editor, {
-            getChanges() {
-              return getBlocksChangedByTransaction(
-                transaction,
-                appendedTransactions,
-              );
-            },
-          });
+          this.emit("onChange", { editor, transaction, appendedTransactions });
         },
       );
       editor._tiptapEditor.on("selectionUpdate", ({ transaction }) => {
@@ -64,20 +65,41 @@ export class EventManager<Editor extends BlockNoteEditor> extends EventEmitter<{
    */
   public onChange(
     callback: (
-      editor: Editor,
+      editor: BlockNoteEditor<BSchema, I, S>,
       ctx: {
-        getChanges(): BlocksChanged<
-          Editor["schema"]["blockSchema"],
-          Editor["schema"]["inlineContentSchema"],
-          Editor["schema"]["styleSchema"]
-        >;
+        getChanges(): BlocksChanged<BSchema, I, S>;
       },
     ) => void,
+    /**
+     * If true, the callback will be triggered when the changes are caused by a remote user
+     * @default true
+     */
+    includeUpdatesFromRemote = true,
   ): Unsubscribe {
-    this.on("onChange", callback);
+    const cb = ({
+      transaction,
+      appendedTransactions,
+    }: {
+      transaction: Transaction;
+      appendedTransactions: Transaction[];
+    }) => {
+      if (!includeUpdatesFromRemote && isRemoteTransaction(transaction)) {
+        // don't trigger the callback if the changes are caused by a remote user
+        return;
+      }
+      callback(this.editor, {
+        getChanges() {
+          return getBlocksChangedByTransaction(
+            transaction,
+            appendedTransactions,
+          );
+        },
+      });
+    };
+    this.on("onChange", cb);
 
     return () => {
-      this.off("onChange", callback);
+      this.off("onChange", cb);
     };
   }
 
@@ -85,7 +107,7 @@ export class EventManager<Editor extends BlockNoteEditor> extends EventEmitter<{
    * Register a callback that will be called when the selection changes.
    */
   public onSelectionChange(
-    callback: (editor: Editor) => void,
+    callback: (editor: BlockNoteEditor<BSchema, I, S>) => void,
     /**
      * If true, the callback will be triggered when the selection changes due to a yjs sync (i.e.: other user was typing)
      */
@@ -93,11 +115,10 @@ export class EventManager<Editor extends BlockNoteEditor> extends EventEmitter<{
   ): Unsubscribe {
     const cb = (e: { transaction: Transaction }) => {
       if (
-        e.transaction.getMeta("$y-sync") &&
-        !includeSelectionChangedByRemote
+        !includeSelectionChangedByRemote &&
+        isRemoteTransaction(e.transaction)
       ) {
-        // selection changed because of a yjs sync (i.e.: other user was typing)
-        // we don't want to trigger the callback in this case
+        // don't trigger the callback if the selection changed because of a remote user
         return;
       }
       callback(this.editor);
@@ -113,7 +134,9 @@ export class EventManager<Editor extends BlockNoteEditor> extends EventEmitter<{
   /**
    * Register a callback that will be called when the editor is mounted.
    */
-  public onMount(callback: (ctx: { editor: Editor }) => void): Unsubscribe {
+  public onMount(
+    callback: (ctx: { editor: BlockNoteEditor<BSchema, I, S> }) => void,
+  ): Unsubscribe {
     this.on("onMount", callback);
 
     return () => {
@@ -124,7 +147,9 @@ export class EventManager<Editor extends BlockNoteEditor> extends EventEmitter<{
   /**
    * Register a callback that will be called when the editor is unmounted.
    */
-  public onUnmount(callback: (ctx: { editor: Editor }) => void): Unsubscribe {
+  public onUnmount(
+    callback: (ctx: { editor: BlockNoteEditor<BSchema, I, S> }) => void,
+  ): Unsubscribe {
     this.on("onUnmount", callback);
 
     return () => {
@@ -132,3 +157,7 @@ export class EventManager<Editor extends BlockNoteEditor> extends EventEmitter<{
     };
   }
 }
+
+function isRemoteTransaction(transaction: Transaction): boolean {
+  return !!transaction.getMeta("$y-sync");
+}
diff --git a/packages/core/src/editor/managers/ExtensionManager.ts b/packages/core/src/editor/managers/ExtensionManager.ts
deleted file mode 100644
index 4d35b68c17..0000000000
--- a/packages/core/src/editor/managers/ExtensionManager.ts
+++ /dev/null
@@ -1,130 +0,0 @@
-import { FilePanelProsemirrorPlugin } from "../../extensions/FilePanel/FilePanelPlugin.js";
-import { FormattingToolbarProsemirrorPlugin } from "../../extensions/FormattingToolbar/FormattingToolbarPlugin.js";
-import { LinkToolbarProsemirrorPlugin } from "../../extensions/LinkToolbar/LinkToolbarPlugin.js";
-import { ShowSelectionPlugin } from "../../extensions/ShowSelection/ShowSelectionPlugin.js";
-import { SideMenuProsemirrorPlugin } from "../../extensions/SideMenu/SideMenuPlugin.js";
-import { SuggestionMenuProseMirrorPlugin } from "../../extensions/SuggestionMenu/SuggestionPlugin.js";
-import { TableHandlesProsemirrorPlugin } from "../../extensions/TableHandles/TableHandlesPlugin.js";
-import { BlockNoteExtension } from "../BlockNoteExtension.js";
-import { BlockNoteEditor } from "../BlockNoteEditor.js";
-
-export class ExtensionManager {
-  constructor(private editor: BlockNoteEditor) {}
-
-  /**
-   * Shorthand to get a typed extension from the editor, by
-   * just passing in the extension class.
-   *
-   * @param ext - The extension class to get
-   * @param key - optional, the key of the extension in the extensions object (defaults to the extension name)
-   * @returns The extension instance
-   */
-  public extension<T extends BlockNoteExtension>(
-    ext: { new (...args: any[]): T } & typeof BlockNoteExtension,
-    key = ext.key(),
-  ): T {
-    const extension = this.editor.extensions[key] as T;
-    if (!extension) {
-      throw new Error(`Extension ${key} not found`);
-    }
-    return extension;
-  }
-
-  /**
-   * Get all extensions
-   */
-  public getExtensions() {
-    return this.editor.extensions;
-  }
-
-  /**
-   * Get a specific extension by key
-   */
-  public getExtension(key: string) {
-    return this.editor.extensions[key];
-  }
-
-  /**
-   * Check if an extension exists
-   */
-  public hasExtension(key: string): boolean {
-    return key in this.editor.extensions;
-  }
-
-  // Plugin getters - these provide access to the core BlockNote plugins
-
-  /**
-   * Get the formatting toolbar plugin
-   */
-  public get formattingToolbar(): FormattingToolbarProsemirrorPlugin {
-    return this.editor.extensions[
-      "formattingToolbar"
-    ] as FormattingToolbarProsemirrorPlugin;
-  }
-
-  /**
-   * Get the link toolbar plugin
-   */
-  public get linkToolbar(): LinkToolbarProsemirrorPlugin<any, any, any> {
-    return this.editor.extensions[
-      "linkToolbar"
-    ] as LinkToolbarProsemirrorPlugin<any, any, any>;
-  }
-
-  /**
-   * Get the side menu plugin
-   */
-  public get sideMenu(): SideMenuProsemirrorPlugin<any, any, any> {
-    return this.editor.extensions["sideMenu"] as SideMenuProsemirrorPlugin<
-      any,
-      any,
-      any
-    >;
-  }
-
-  /**
-   * Get the suggestion menus plugin
-   */
-  public get suggestionMenus(): SuggestionMenuProseMirrorPlugin<any, any, any> {
-    return this.editor.extensions[
-      "suggestionMenus"
-    ] as SuggestionMenuProseMirrorPlugin<any, any, any>;
-  }
-
-  /**
-   * Get the file panel plugin (if available)
-   */
-  public get filePanel(): FilePanelProsemirrorPlugin<any, any> | undefined {
-    return this.editor.extensions["filePanel"] as
-      | FilePanelProsemirrorPlugin<any, any>
-      | undefined;
-  }
-
-  /**
-   * Get the table handles plugin (if available)
-   */
-  public get tableHandles():
-    | TableHandlesProsemirrorPlugin<any, any>
-    | undefined {
-    return this.editor.extensions["tableHandles"] as
-      | TableHandlesProsemirrorPlugin<any, any>
-      | undefined;
-  }
-
-  /**
-   * Get the show selection plugin
-   */
-  public get showSelectionPlugin(): ShowSelectionPlugin {
-    return this.editor.extensions["showSelection"] as ShowSelectionPlugin;
-  }
-
-  /**
-   * Check if collaboration is enabled (Yjs or Liveblocks)
-   */
-  public get isCollaborationEnabled(): boolean {
-    return (
-      this.hasExtension("ySyncPlugin") ||
-      this.hasExtension("liveblocksExtension")
-    );
-  }
-}
diff --git a/packages/core/src/editor/managers/ExtensionManager/extensions.ts b/packages/core/src/editor/managers/ExtensionManager/extensions.ts
new file mode 100644
index 0000000000..e4e8253853
--- /dev/null
+++ b/packages/core/src/editor/managers/ExtensionManager/extensions.ts
@@ -0,0 +1,152 @@
+import {
+  AnyExtension as AnyTiptapExtension,
+  extensions,
+  Node,
+  Extension as TiptapExtension,
+} from "@tiptap/core";
+import { Gapcursor } from "@tiptap/extension-gapcursor";
+import { Link } from "@tiptap/extension-link";
+import { Text } from "@tiptap/extension-text";
+import { createDropFileExtension } from "../../../api/clipboard/fromClipboard/fileDropExtension.js";
+import { createPasteFromClipboardExtension } from "../../../api/clipboard/fromClipboard/pasteExtension.js";
+import { createCopyToClipboardExtension } from "../../../api/clipboard/toClipboard/copyExtension.js";
+import {
+  BackgroundColorExtension,
+  HardBreak,
+  KeyboardShortcutsExtension,
+  TextAlignmentExtension,
+  TextColorExtension,
+  UniqueID,
+  SuggestionAddMark,
+  SuggestionDeleteMark,
+  SuggestionModificationMark,
+} from "../../../extensions/tiptap-extensions/index.js";
+import {
+  DEFAULT_LINK_PROTOCOL,
+  VALID_LINK_PROTOCOLS,
+} from "../../../extensions/LinkToolbar/protocols.js";
+import { BlockContainer, BlockGroup, Doc } from "../../../pm-nodes/index.js";
+import {
+  BlockNoteEditor,
+  BlockNoteEditorOptions,
+} from "../../BlockNoteEditor.js";
+import { SuggestionMenu } from "../../../extensions/index.js";
+
+// TODO remove linkify completely by vendoring the link extension & dropping linkifyjs as a dependency
+let LINKIFY_INITIALIZED = false;
+
+/**
+ * Get all the Tiptap extensions BlockNote is configured with by default
+ */
+export function getDefaultTiptapExtensions(
+  editor: BlockNoteEditor<any, any, any>,
+  options: BlockNoteEditorOptions<any, any, any>,
+) {
+  const tiptapExtensions: AnyTiptapExtension[] = [
+    extensions.ClipboardTextSerializer,
+    extensions.Commands,
+    extensions.Editable,
+    extensions.FocusEvents,
+    extensions.Tabindex,
+    Gapcursor,
+
+    UniqueID.configure({
+      // everything from bnBlock group (nodes that represent a BlockNote block should have an id)
+      types: ["blockContainer", "columnList", "column"],
+      setIdAttribute: options.setIdAttribute,
+    }),
+    HardBreak,
+    Text,
+
+    // marks:
+    SuggestionAddMark,
+    SuggestionDeleteMark,
+    SuggestionModificationMark,
+    Link.extend({
+      inclusive: false,
+    }).configure({
+      defaultProtocol: DEFAULT_LINK_PROTOCOL,
+      // only call this once if we have multiple editors installed. Or fix https://github.com/ueberdosis/tiptap/issues/5450
+      protocols: LINKIFY_INITIALIZED ? [] : VALID_LINK_PROTOCOLS,
+    }),
+    ...(Object.values(editor.schema.styleSpecs).map((styleSpec) => {
+      return styleSpec.implementation.mark.configure({
+        editor: editor,
+      });
+    }) as any[]),
+
+    TextColorExtension,
+
+    BackgroundColorExtension,
+    TextAlignmentExtension,
+
+    // make sure escape blurs editor, so that we can tab to other elements in the host page (accessibility)
+    TiptapExtension.create({
+      name: "OverrideEscape",
+      addKeyboardShortcuts: () => {
+        return {
+          Escape: () => {
+            // TODO should this be like this?
+            if (editor.getExtension(SuggestionMenu)?.shown()) {
+              // escape is handled by suggestionmenu
+              return false;
+            }
+            editor.blur();
+            return true;
+          },
+        };
+      },
+    }),
+
+    // nodes
+    Doc,
+    BlockContainer.configure({
+      editor: editor,
+      domAttributes: options.domAttributes,
+    }),
+    KeyboardShortcutsExtension.configure({
+      editor: editor,
+      tabBehavior: options.tabBehavior,
+    }),
+    BlockGroup.configure({
+      domAttributes: options.domAttributes,
+    }),
+    ...Object.values(editor.schema.inlineContentSpecs)
+      .filter((a) => a.config !== "link" && a.config !== "text")
+      .map((inlineContentSpec) => {
+        return inlineContentSpec.implementation!.node.configure({
+          editor: editor,
+        });
+      }),
+
+    ...Object.values(editor.schema.blockSpecs).flatMap((blockSpec) => {
+      return [
+        // the node extension implementations
+        ...("node" in blockSpec.implementation
+          ? [
+              (blockSpec.implementation.node as Node).configure({
+                editor: editor,
+                domAttributes: options.domAttributes,
+              }),
+            ]
+          : []),
+      ];
+    }),
+    createCopyToClipboardExtension(editor),
+    createPasteFromClipboardExtension(
+      editor,
+      options.pasteHandler ||
+        ((context: {
+          defaultPasteHandler: (context?: {
+            prioritizeMarkdownOverHTML?: boolean;
+            plainTextAsMarkdown?: boolean;
+          }) => boolean | undefined;
+        }) => context.defaultPasteHandler()),
+    ),
+    createDropFileExtension(editor),
+  ];
+
+  LINKIFY_INITIALIZED = true;
+
+  return tiptapExtensions;
+}
diff --git a/packages/core/src/editor/managers/ExtensionManager/index.ts b/packages/core/src/editor/managers/ExtensionManager/index.ts
new file mode 100644
index 0000000000..dadce1e5bc
--- /dev/null
+++ b/packages/core/src/editor/managers/ExtensionManager/index.ts
@@ -0,0 +1,456 @@
+import { InputRule, inputRules } from "@handlewithcare/prosemirror-inputrules";
+import {
+  AnyExtension as AnyTiptapExtension,
+  Extension as TiptapExtension,
+} from "@tiptap/core";
+import { keymap } from "@tiptap/pm/keymap";
+import { Plugin } from "prosemirror-state";
+import { updateBlockTr } from "../../../api/blockManipulation/commands/updateBlock/updateBlock.js";
+import { getBlockInfoFromTransaction } from "../../../api/getBlockInfoFromPos.js";
+import { DEFAULT_EXTENSIONS } from "../../../extensions/index.js";
+import { sortByDependencies } from "../../../util/topo-sort.js";
+import type {
+  BlockNoteEditor,
+  BlockNoteEditorOptions,
+} from "../../BlockNoteEditor.js";
+import type { Extension, ExtensionFactory } from "../../BlockNoteExtension.js";
+import { getDefaultTiptapExtensions } from "./extensions.js";
+
+export class ExtensionManager {
+  /**
+   * A set of extension keys which are disabled by the options
+   */
+  private disabledExtensions = new Set<string>();
+  /**
+   * A list of all the extensions that are registered to the editor
+   */
+  private extensions: Extension[] = [];
+  /**
+   * A map of all the abort controllers for each extension that has an init method defined
+   */
+  private abortMap = new Map<Extension, AbortController>();
+  /**
+   * A map of all the extension factories that are registered to the editor
+   */
+  private extensionFactories = new Map<ExtensionFactory, Extension>();
+  /**
+   * Because a single blocknote extension can both have it's own prosemirror plugins & additional generated ones (e.g. keymap & input rules plugins)
+   * We need to keep track of all the plugins for each extension, so that we can remove them when the extension is unregistered
+   */
+  private extensionPlugins: Map<Extension, Plugin[]> = new Map();
+
+  constructor(
+    private editor: BlockNoteEditor<any, any, any>,
+    private options: BlockNoteEditorOptions<any, any, any>,
+  ) {
+    /**
+     * When the editor is first mounted, we need to initialize all the extensions
+     */
+    editor.onMount(() => {
+      for (const extension of this.extensions) {
+        // If the extension has an init function, we can initialize it, otherwise, it is already added to the editor
+        if (extension.mount) {
+          // We create an abort controller for each extension, so that we can abort the extension when the editor is unmounted
+          const abortController = new AbortController();
+          const unmountCallback = extension.mount({
+            dom: editor.prosemirrorView.dom,
+            root: editor.prosemirrorView.root,
+            signal: abortController.signal,
+          });
+          // If the extension returns a method to unmount it, we can register it to be called when the abort controller is aborted
+          if (unmountCallback) {
+            abortController.signal.addEventListener("abort", () => {
+              unmountCallback();
+            });
+          }
+          // Keep track of the abort controller for each extension, so that we can abort it when the editor is unmounted
+          this.abortMap.set(extension, abortController);
+        }
+      }
+    });
+
+    /**
+     * When the editor is unmounted, we need to abort all the extensions' abort controllers
+     */
+    editor.onUnmount(() => {
+      for (const [extension, abortController] of this.abortMap.entries()) {
+        // No longer track the abort controller for this extension
+        this.abortMap.delete(extension);
+        // Abort each extension's abort controller
+        abortController.abort();
+      }
+    });
+
+    // TODO do disabled extensions need to be only for editor base extensions? Or all of them?
+    this.disabledExtensions = new Set(options.disableExtensions || []);
+
+    // Add the default extensions
+    for (const extension of DEFAULT_EXTENSIONS) {
+      this.addExtension(extension);
+    }
+
+    // Add the extensions from the options
+    for (const extension of this.options.extensions ?? []) {
+      this.addExtension(extension);
+    }
+
+    // Add the extensions from blocks specs
+    for (const block of Object.values(this.editor.schema.blockSpecs)) {
+      for (const extension of block.extensions ?? []) {
+        this.addExtension(extension);
+      }
+    }
+  }
+
+  /**
+   * Register one or more extensions to the editor after the editor is initialized.
+   *
+   * This allows users to switch on & off extensions "at runtime".
+   */
+  public registerExtension(
+    extension:
+      | undefined
+      | Extension
+      | ExtensionFactory
+      | (Extension | ExtensionFactory | undefined)[],
+  ): void {
+    const extensions = ([] as (Extension | ExtensionFactory | undefined)[])
+      .concat(extension)
+      .filter(Boolean) as (Extension | ExtensionFactory)[];
+
+    if (!extensions.length) {
+      // eslint-disable-next-line no-console
+      console.warn(`No extensions found to register`, extension);
+      return;
+    }
+
+    const registeredExtensions = extensions
+      .map((extension) => this.addExtension(extension))
+      .filter(Boolean) as Extension[];
+
+    const pluginsToAdd = new Set<Plugin>();
+    for (const extension of registeredExtensions) {
+      if (extension?.tiptapExtensions) {
+        // This is necessary because this can only switch out prosemirror plugins at runtime,
+        // it can't switch out Tiptap extensions since that can have more widespread effects (since a Tiptap extension can even add/remove to the schema).
+
+        // eslint-disable-next-line no-console
+        console.warn(
+          `Extension ${extension.key} has tiptap extensions, but these cannot be changed after initializing the editor. Please separate the extension into multiple extensions if you want to add them, or re-initialize the editor.`,
+          extension,
+        );
+      }
+
+      this.getProsemirrorPluginsFromExtension(extension).forEach((plugin) => {
+        pluginsToAdd.add(plugin);
+      });
+    }
+
+    // TODO there isn't a great way to do sorting right now. This is something that should be improved in the future.
+    this.updatePlugins((plugins) => [...plugins, ...pluginsToAdd]);
+  }
+
+  /**
+   * Register an extension to the editor
+   * @param extension - The extension to register
+   * @returns The extension instance
+   */
+  private addExtension<T extends ExtensionFactory | Extension>(
+    extension: T,
+  ): T extends ExtensionFactory
+    ? ReturnType<T>
+    : T extends Extension | undefined
+      ? T
+      : never {
+    let instance: Extension | undefined;
+    if (typeof extension === "function") {
+      instance = extension(this.editor, this.options);
+    } else {
+      instance = extension;
+    }
+
+    if (!instance || this.disabledExtensions.has(instance.key)) {
+      return undefined as any;
+    }
+
+    // Now that we know that the extension is not disabled, we can add it to the extension factories
+    if (typeof extension === "function") {
+      this.extensionFactories.set(extension, instance);
+    }
+
+    this.extensions.push(instance);
+
+    return instance as any;
+  }
+
+  /**
+   * Resolve an extension or a list of extensions into a list of extension instances
+   * @param toResolve - The extension or list of extensions to resolve
+   * @returns A list of extension instances
+   */
+  private resolveExtensions(
+    toResolve:
+      | undefined
+      | string
+      | Extension
+      | ExtensionFactory
+      | (Extension | ExtensionFactory | string | undefined)[],
+  ): Extension[] {
+    const extensions = [] as Extension[];
+    if (typeof toResolve === "function") {
+      const instance = this.extensionFactories.get(toResolve);
+      if (instance) {
+        extensions.push(instance);
+      }
+    } else if (Array.isArray(toResolve)) {
+      for (const extension of toResolve) {
+        extensions.push(...this.resolveExtensions(extension));
+      }
+    } else if (typeof toResolve === "object" && "key" in toResolve) {
+      extensions.push(toResolve);
+    } else if (typeof toResolve === "string") {
+      const instance = this.extensions.find((e) => e.key === toResolve);
+      if (instance) {
+        extensions.push(instance);
+      }
+    }
+    return extensions;
+  }
+
+  /**
+   * Unregister an extension from the editor
+   * @param toUnregister - The extension to unregister
+   * @returns void
+   */
+  public unregisterExtension(
+    toUnregister:
+      | undefined
+      | string
+      | Extension
+      | ExtensionFactory
+      | (Extension | ExtensionFactory | string | undefined)[],
+  ): void {
+    const extensions = this.resolveExtensions(toUnregister);
+
+    if (!extensions.length) {
+      // eslint-disable-next-line no-console
+      console.warn(`No extensions found to unregister`, toUnregister);
+      return;
+    }
+    let didWarn = false;
+
+    const pluginsToRemove = new Set<Plugin>();
+    for (const extension of extensions) {
+      this.extensions = this.extensions.filter((e) => e !== extension);
+      this.extensionFactories.entries().forEach(([factory, instance]) => {
+        if (instance === extension) {
+          this.extensionFactories.delete(factory);
+        }
+      });
+      this.abortMap.get(extension)?.abort();
+      this.abortMap.delete(extension);
+
+      const plugins = this.extensionPlugins.get(extension);
+      plugins?.forEach((plugin) => {
+        pluginsToRemove.add(plugin);
+      });
+      this.extensionPlugins.delete(extension);
+
+      if (extension.tiptapExtensions && !didWarn) {
+        didWarn = true;
+        // eslint-disable-next-line no-console
+        console.warn(
+          `Extension ${extension.key} has tiptap extensions, but they will not be removed. Please separate the extension into multiple extensions if you want to remove them, or re-initialize the editor.`,
+          toUnregister,
+        );
+      }
+    }
+
+    this.updatePlugins((plugins) =>
+      plugins.filter((plugin) => !pluginsToRemove.has(plugin)),
+    );
+  }
+
+  /**
+   * Allows resetting the current prosemirror state's plugins
+   * @param update - A function that takes the current plugins and returns the new plugins
+   * @returns void
+   */
+  private updatePlugins(update: (plugins: Plugin[]) => Plugin[]): void {
+    const currentState = this.editor.prosemirrorState;
+
+    const state = currentState.reconfigure({
+      plugins: update(currentState.plugins.slice()),
+    });
+
+    this.editor.prosemirrorView.updateState(state);
+  }
+
+  /**
+   * Get all the extensions that are registered to the editor
+   */
+  public getTiptapExtensions(): AnyTiptapExtension[] {
+    // Start with the default tiptap extensions
+    const tiptapExtensions = getDefaultTiptapExtensions(
+      this.editor,
+      this.options,
+    );
+    // TODO filter out the default extensions via the disabledExtensions set?
+
+    const getPriority = sortByDependencies(this.extensions);
+
+    for (const extension of this.extensions) {
+      if (extension.tiptapExtensions) {
+        tiptapExtensions.push(...extension.tiptapExtensions);
+      }
+
+      const prosemirrorPlugins =
+        this.getProsemirrorPluginsFromExtension(extension);
+      // Sometimes a blocknote extension might need to make additional prosemirror plugins, so we generate them here
+      if (prosemirrorPlugins.length) {
+        tiptapExtensions.push(
+          TiptapExtension.create({
+            name: extension.key,
+            priority: getPriority(extension.key),
+            addProseMirrorPlugins: () => prosemirrorPlugins,
+          }),
+        );
+      }
+    }
+
+    // Add any tiptap extensions from the `_tiptapOptions`
+    for (const extension of this.options._tiptapOptions?.extensions ?? []) {
+      tiptapExtensions.push(extension);
+    }
+
+    return tiptapExtensions;
+  }
+
+  /**
+   * This maps a blocknote extension into an array of Prosemirror plugins if it has any of the following:
+   * - plugins
+   * - keyboard shortcuts
+   * - input rules
+   */
+  private getProsemirrorPluginsFromExtension(extension: Extension): Plugin[] {
+    if (
+      !extension.prosemirrorPlugins?.length &&
+      !extension.keyboardShortcuts?.length &&
+      !extension.inputRules?.length
+    ) {
+      // We can bail out early if the extension has no features to add to the tiptap editor
+      return [];
+    }
+
+    const plugins: Plugin[] = [...(extension.prosemirrorPlugins ?? [])];
+
+    this.extensionPlugins.set(extension, plugins);
+
+    if (extension.inputRules?.length) {
+      plugins.push(
+        inputRules({
+          rules: extension.inputRules.map((inputRule) => {
+            return new InputRule(inputRule.find, (state, match, start, end) => {
+              const replaceWith = inputRule.replace({
+                match,
+                range: { from: start, to: end },
+                editor: this.editor,
+              });
+              if (replaceWith) {
+                const cursorPosition = this.editor.getTextCursorPosition();
+
+                if (
+                  this.editor.schema.blockSchema[cursorPosition.block.type]
+                    .content !== "inline"
+                ) {
+                  return null;
+                }
+
+                const blockInfo = getBlockInfoFromTransaction(state.tr);
+                const tr = state.tr.deleteRange(start, end);
+
+                updateBlockTr(tr, blockInfo.bnBlock.beforePos, replaceWith);
+                return tr;
+              }
+              return null;
+            });
+          }),
+        }),
+      );
+    }
+
+    if (extension.keyboardShortcuts?.length) {
+      plugins.push(
+        keymap(
+          Object.fromEntries(
+            Object.entries(extension.keyboardShortcuts).map(([key, value]) => [
+              key,
+              () => value({ editor: this.editor }),
+            ]),
+          ),
+        ),
+      );
+    }
+
+    return plugins;
+  }
+
+  /**
+   * Get all extensions
+   */
+  public getExtensions(): Map<string, Extension> {
+    return new Map(
+      this.extensions.map((extension) => [extension.key, extension]),
+    );
+  }
+
+  /**
+   * Get a specific extension by it's instance
+   */
+  public getExtension<
+    T extends ExtensionFactory | Extension | string | undefined,
+  >(
+    extension: T,
+  ):
+    | (T extends ExtensionFactory
+        ? ReturnType<T>
+        : T extends Extension
+          ? T
+          : T extends string
+            ? Extension
+            : never)
+    | undefined {
+    if (!extension) {
+      return undefined;
+    } else if (typeof extension === "string") {
+      const instance = this.extensions.find((e) => e.key === extension);
+      if (!instance) {
+        return undefined;
+      }
+      return instance as any;
+    } else if (typeof extension === "function") {
+      const instance = this.extensionFactories.get(extension);
+      if (!instance) {
+        return undefined;
+      }
+      return instance as any;
+    } else if (typeof extension === "object" && "key" in extension) {
+      return this.getExtension(extension.key) as any;
+    }
+    throw new Error(`Invalid extension type: ${typeof extension}`);
+  }
+
+  /**
+   * Check if an extension exists
+   */
+  public hasExtension(key: string | Extension | ExtensionFactory): boolean {
+    if (typeof key === "string") {
+      return this.extensions.some((e) => e.key === key);
+    } else if (typeof key === "object" && "key" in key) {
+      return this.extensions.some((e) => e.key === key.key);
+    } else if (typeof key === "function") {
+      return this.extensionFactories.has(key);
+    }
+    return false;
+  }
+}
diff --git a/packages/core/src/editor/managers/SelectionManager.ts b/packages/core/src/editor/managers/SelectionManager.ts
index a91487251c..6fb16b7c6a 100644
--- a/packages/core/src/editor/managers/SelectionManager.ts
+++ b/packages/core/src/editor/managers/SelectionManager.ts
@@ -109,6 +109,10 @@ export class SelectionManager<
       }
     }
 
-    return posToDOMRect(this.editor.prosemirrorView, from, to);
+    return posToDOMRect(
+      this.editor.prosemirrorView,
+      from,
+      to,
+    ).toJSON() as DOMRect;
   }
 }
diff --git a/packages/core/src/editor/managers/StateManager.ts b/packages/core/src/editor/managers/StateManager.ts
index 6e9f2be42b..31da76e6f7 100644
--- a/packages/core/src/editor/managers/StateManager.ts
+++ b/packages/core/src/editor/managers/StateManager.ts
@@ -1,21 +1,10 @@
-import { redo, undo } from "@tiptap/pm/history";
 import { Command, Transaction } from "prosemirror-state";
+import type { YUndo } from "../../extensions/Collaboration/YUndo.js";
+import type { History } from "../../extensions/History/History.js";
 import { BlockNoteEditor } from "../BlockNoteEditor.js";
 
 export class StateManager {
-  constructor(
-    private editor: BlockNoteEditor,
-    private options?: {
-      /**
-       * Swap the default undo command with a custom command.
-       */
-      undo?: typeof undo;
-      /**
-       * Swap the default redo command with a custom command.
-       */
-      redo?: typeof redo;
-    },
-  ) {}
+  constructor(private editor: BlockNoteEditor<any, any, any>) {}
 
   /**
    * Stores the currently active transaction, which is the accumulated transaction from all {@link dispatch} calls during a {@link transact} calls
@@ -225,14 +214,43 @@ export class StateManager {
   /**
    * Undo the last action.
    */
-  public undo() {
-    return this.exec(this.options?.undo ?? undo);
+  public undo(): boolean {
+    // Purposefully not using the UndoPlugin to not import y-prosemirror when not needed
+    const undoPlugin = this.editor.getExtension("yUndo") as ReturnType<
+      typeof YUndo
+    >;
+    if (undoPlugin) {
+      return this.exec(undoPlugin.undoCommand);
+    }
+
+    const historyPlugin = this.editor.getExtension("history") as ReturnType<
+      typeof History
+    >;
+    if (historyPlugin) {
+      return this.exec(historyPlugin.undoCommand);
+    }
+
+    throw new Error("No undo plugin found");
   }
 
   /**
    * Redo the last action.
    */
   public redo() {
-    return this.exec(this.options?.redo ?? redo);
+    const undoPlugin = this.editor.getExtension("yUndo") as ReturnType<
+      typeof YUndo
+    >;
+    if (undoPlugin) {
+      return this.exec(undoPlugin.redoCommand);
+    }
+
+    const historyPlugin = this.editor.getExtension("history") as ReturnType<
+      typeof History
+    >;
+    if (historyPlugin) {
+      return this.exec(historyPlugin.redoCommand);
+    }
+
+    throw new Error("No redo plugin found");
   }
 }
diff --git a/packages/core/src/editor/managers/index.ts b/packages/core/src/editor/managers/index.ts
index c4f64b6e3d..3a23d009db 100644
--- a/packages/core/src/editor/managers/index.ts
+++ b/packages/core/src/editor/managers/index.ts
@@ -1,11 +1,8 @@
 export { BlockManager } from "./BlockManager.js";
-export {
-  CollaborationManager,
-  type CollaborationOptions,
-} from "./CollaborationManager.js";
+export { CollaborationManager } from "./CollaborationManager.js";
 export { EventManager } from "./EventManager.js";
 export { ExportManager } from "./ExportManager.js";
-export { ExtensionManager } from "./ExtensionManager.js";
+export { ExtensionManager } from "./ExtensionManager/index.js";
 export { SelectionManager } from "./SelectionManager.js";
 export { StateManager } from "./StateManager.js";
 export { StyleManager } from "./StyleManager.js";
diff --git a/packages/core/src/extensions/BlockChange/BlockChangePlugin.ts b/packages/core/src/extensions/BlockChange/BlockChange.ts
similarity index 61%
rename from packages/core/src/extensions/BlockChange/BlockChangePlugin.ts
rename to packages/core/src/extensions/BlockChange/BlockChange.ts
index 94dfdac263..3f054e519f 100644
--- a/packages/core/src/extensions/BlockChange/BlockChangePlugin.ts
+++ b/packages/core/src/extensions/BlockChange/BlockChange.ts
@@ -3,25 +3,19 @@ import {
   BlocksChanged,
   getBlocksChangedByTransaction,
 } from "../../api/getBlocksChangedByTransaction.js";
-import { BlockNoteExtension } from "../../editor/BlockNoteExtension.js";
+import { createExtension } from "../../editor/BlockNoteExtension.js";
 
 /**
  * This plugin can filter transactions before they are applied to the editor, but with a higher-level API than `filterTransaction` from prosemirror.
  */
-export class BlockChangePlugin extends BlockNoteExtension {
-  public static key() {
-    return "blockChange";
-  }
-
-  private beforeChangeCallbacks: ((context: {
+export const BlockChange = createExtension((_editor, _options) => {
+  const beforeChangeCallbacks: ((context: {
     getChanges: () => BlocksChanged<any, any, any>;
     tr: Transaction;
   }) => boolean | void)[] = [];
-
-  constructor() {
-    super();
-
-    this.addProsemirrorPlugin(
+  return {
+    key: "blockChange",
+    prosemirrorPlugins: [
       new Plugin({
         key: new PluginKey("blockChange"),
         filterTransaction: (tr) => {
@@ -29,7 +23,7 @@ export class BlockChangePlugin extends BlockNoteExtension {
             | ReturnType<typeof getBlocksChangedByTransaction>
             | undefined = undefined;
 
-          return this.beforeChangeCallbacks.reduce((acc, cb) => {
+          return beforeChangeCallbacks.reduce((acc, cb) => {
             if (acc === false) {
               // We only care that we hit a `false` result, so we can stop iterating.
               return acc;
@@ -49,21 +43,25 @@ export class BlockChangePlugin extends BlockNoteExtension {
           }, true);
         },
       }),
-    );
-  }
+    ],
 
-  public subscribe(
-    callback: (context: {
-      getChanges: () => BlocksChanged<any, any, any>;
-      tr: Transaction;
-    }) => boolean | void,
-  ) {
-    this.beforeChangeCallbacks.push(callback);
+    /**
+     * Subscribe to the block change events.
+     */
+    subscribe(
+      callback: (context: {
+        getChanges: () => BlocksChanged<any, any, any>;
+        tr: Transaction;
+      }) => boolean | void,
+    ) {
+      beforeChangeCallbacks.push(callback);
 
-    return () => {
-      this.beforeChangeCallbacks = this.beforeChangeCallbacks.filter(
-        (cb) => cb !== callback,
-      );
-    };
-  }
-}
+      return () => {
+        beforeChangeCallbacks.splice(
+          beforeChangeCallbacks.indexOf(callback),
+          1,
+        );
+      };
+    },
+  } as const;
+});
diff --git a/packages/core/src/extensions/Collaboration/CursorPlugin.ts b/packages/core/src/extensions/Collaboration/CursorPlugin.ts
deleted file mode 100644
index 5e2a45d116..0000000000
--- a/packages/core/src/extensions/Collaboration/CursorPlugin.ts
+++ /dev/null
@@ -1,189 +0,0 @@
-import { defaultSelectionBuilder, yCursorPlugin } from "y-prosemirror";
-import { Awareness } from "y-protocols/awareness.js";
-import * as Y from "yjs";
-import { BlockNoteExtension } from "../../editor/BlockNoteExtension.js";
-
-export type CollaborationUser = {
-  name: string;
-  color: string;
-  [key: string]: string;
-};
-
-export class CursorPlugin extends BlockNoteExtension {
-  public static key() {
-    return "yCursorPlugin";
-  }
-
-  private provider: { awareness: Awareness };
-  private recentlyUpdatedCursors: Map<
-    number,
-    { element: HTMLElement; hideTimeout: NodeJS.Timeout | undefined }
-  >;
-  constructor(
-    private collaboration: {
-      fragment: Y.XmlFragment;
-      user: CollaborationUser;
-      provider: { awareness: Awareness };
-      renderCursor?: (user: CollaborationUser) => HTMLElement;
-      showCursorLabels?: "always" | "activity";
-    },
-  ) {
-    super();
-    this.provider = collaboration.provider;
-    this.recentlyUpdatedCursors = new Map();
-
-    this.provider.awareness.setLocalStateField("user", collaboration.user);
-
-    if (collaboration.showCursorLabels !== "always") {
-      this.provider.awareness.on(
-        "change",
-        ({
-          updated,
-        }: {
-          added: Array<number>;
-          updated: Array<number>;
-          removed: Array<number>;
-        }) => {
-          for (const clientID of updated) {
-            const cursor = this.recentlyUpdatedCursors.get(clientID);
-
-            if (cursor) {
-              cursor.element.setAttribute("data-active", "");
-
-              if (cursor.hideTimeout) {
-                clearTimeout(cursor.hideTimeout);
-              }
-
-              this.recentlyUpdatedCursors.set(clientID, {
-                element: cursor.element,
-                hideTimeout: setTimeout(() => {
-                  cursor.element.removeAttribute("data-active");
-                }, 2000),
-              });
-            }
-          }
-        },
-      );
-    }
-
-    this.addProsemirrorPlugin(
-      yCursorPlugin(this.provider.awareness, {
-        selectionBuilder: defaultSelectionBuilder,
-        cursorBuilder: this.renderCursor,
-      }),
-    );
-  }
-
-  public get priority() {
-    return 999;
-  }
-
-  private renderCursor = (user: CollaborationUser, clientID: number) => {
-    let cursorData = this.recentlyUpdatedCursors.get(clientID);
-
-    if (!cursorData) {
-      const cursorElement = (
-        this.collaboration.renderCursor ?? CursorPlugin.defaultCursorRender
-      )(user);
-
-      if (this.collaboration.showCursorLabels !== "always") {
-        cursorElement.addEventListener("mouseenter", () => {
-          const cursor = this.recentlyUpdatedCursors.get(clientID)!;
-          cursor.element.setAttribute("data-active", "");
-
-          if (cursor.hideTimeout) {
-            clearTimeout(cursor.hideTimeout);
-            this.recentlyUpdatedCursors.set(clientID, {
-              element: cursor.element,
-              hideTimeout: undefined,
-            });
-          }
-        });
-
-        cursorElement.addEventListener("mouseleave", () => {
-          const cursor = this.recentlyUpdatedCursors.get(clientID)!;
-
-          this.recentlyUpdatedCursors.set(clientID, {
-            element: cursor.element,
-            hideTimeout: setTimeout(() => {
-              cursor.element.removeAttribute("data-active");
-            }, 2000),
-          });
-        });
-      }
-
-      cursorData = {
-        element: cursorElement,
-        hideTimeout: undefined,
-      };
-
-      this.recentlyUpdatedCursors.set(clientID, cursorData);
-    }
-
-    return cursorData.element;
-  };
-
-  public updateUser = (user: {
-    name: string;
-    color: string;
-    [key: string]: string;
-  }) => {
-    this.provider.awareness.setLocalStateField("user", user);
-  };
-
-  /**
-   * Determine whether the foreground color should be white or black based on a provided background color
-   * Inspired by: https://stackoverflow.com/a/3943023
-   *
-   */
-  public static isDarkColor(bgColor: string): boolean {
-    const color = bgColor.charAt(0) === "#" ? bgColor.substring(1, 7) : bgColor;
-    const r = parseInt(color.substring(0, 2), 16); // hexToR
-    const g = parseInt(color.substring(2, 4), 16); // hexToG
-    const b = parseInt(color.substring(4, 6), 16); // hexToB
-    const uicolors = [r / 255, g / 255, b / 255];
-    const c = uicolors.map((col) => {
-      if (col <= 0.03928) {
-        return col / 12.92;
-      }
-      return Math.pow((col + 0.055) / 1.055, 2.4);
-    });
-    const L = 0.2126 * c[0] + 0.7152 * c[1] + 0.0722 * c[2];
-    return L <= 0.179;
-  }
-
-  public static defaultCursorRender = (user: CollaborationUser) => {
-    const cursorElement = document.createElement("span");
-
-    cursorElement.classList.add("bn-collaboration-cursor__base");
-
-    const caretElement = document.createElement("span");
-    caretElement.setAttribute("contentedEditable", "false");
-    caretElement.classList.add("bn-collaboration-cursor__caret");
-    caretElement.setAttribute(
-      "style",
-      `background-color: ${user.color}; color: ${
-        CursorPlugin.isDarkColor(user.color) ? "white" : "black"
-      }`,
-    );
-
-    const labelElement = document.createElement("span");
-
-    labelElement.classList.add("bn-collaboration-cursor__label");
-    labelElement.setAttribute(
-      "style",
-      `background-color: ${user.color}; color: ${
-        CursorPlugin.isDarkColor(user.color) ? "white" : "black"
-      }`,
-    );
-    labelElement.insertBefore(document.createTextNode(user.name), null);
-
-    caretElement.insertBefore(labelElement, null);
-
-    cursorElement.insertBefore(document.createTextNode("\u2060"), null); // Non-breaking space
-    cursorElement.insertBefore(caretElement, null);
-    cursorElement.insertBefore(document.createTextNode("\u2060"), null); // Non-breaking space
-
-    return cursorElement;
-  };
-}
diff --git a/packages/core/src/extensions/Collaboration/ForkYDocPlugin.test.ts b/packages/core/src/extensions/Collaboration/ForkYDoc.test.ts
similarity index 93%
rename from packages/core/src/extensions/Collaboration/ForkYDocPlugin.test.ts
rename to packages/core/src/extensions/Collaboration/ForkYDoc.test.ts
index eca3ba26f2..a9c16d2ec6 100644
--- a/packages/core/src/extensions/Collaboration/ForkYDocPlugin.test.ts
+++ b/packages/core/src/extensions/Collaboration/ForkYDoc.test.ts
@@ -2,6 +2,7 @@ import { expect, it } from "vitest";
 import * as Y from "yjs";
 import { Awareness } from "y-protocols/awareness";
 import { BlockNoteEditor } from "../../index.js";
+import { ForkYDoc } from "./ForkYDoc.js";
 
 /**
  * @vitest-environment jsdom
@@ -36,7 +37,7 @@ it("can fork a document", async () => {
     "__snapshots__/fork-yjs-snap-editor.json",
   );
 
-  editor.forkYDocPlugin!.fork();
+  editor.getExtension(ForkYDoc)!.fork();
 
   editor.replaceBlocks(editor.document, [
     {
@@ -83,7 +84,7 @@ it("can merge a document", async () => {
     "__snapshots__/fork-yjs-snap-editor.json",
   );
 
-  editor.forkYDocPlugin!.fork();
+  editor.getExtension(ForkYDoc)!.fork();
 
   editor.replaceBlocks(editor.document, [
     {
@@ -99,7 +100,7 @@ it("can merge a document", async () => {
     "__snapshots__/fork-yjs-snap-editor-forked.json",
   );
 
-  editor.forkYDocPlugin!.merge({ keepChanges: false });
+  editor.getExtension(ForkYDoc)!.merge({ keepChanges: false });
 
   await expect(fragment.toJSON()).toMatchFileSnapshot(
     "__snapshots__/fork-yjs-snap.html",
@@ -139,7 +140,7 @@ it("can fork an keep the changes to the original document", async () => {
     "__snapshots__/fork-yjs-snap-editor.json",
   );
 
-  editor.forkYDocPlugin!.fork();
+  editor.getExtension(ForkYDoc)!.fork();
 
   editor.replaceBlocks(editor.document, [
     {
@@ -155,7 +156,7 @@ it("can fork an keep the changes to the original document", async () => {
     "__snapshots__/fork-yjs-snap-editor-forked.json",
   );
 
-  editor.forkYDocPlugin!.merge({ keepChanges: true });
+  editor.getExtension(ForkYDoc)!.merge({ keepChanges: true });
 
   await expect(fragment.toJSON()).toMatchFileSnapshot(
     "__snapshots__/fork-yjs-snap-forked.html",
diff --git a/packages/core/src/extensions/Collaboration/ForkYDoc.ts b/packages/core/src/extensions/Collaboration/ForkYDoc.ts
new file mode 100644
index 0000000000..17ee4dc165
--- /dev/null
+++ b/packages/core/src/extensions/Collaboration/ForkYDoc.ts
@@ -0,0 +1,147 @@
+import { yUndoPluginKey } from "y-prosemirror";
+import * as Y from "yjs";
+import {
+  createExtension,
+  createStore,
+} from "../../editor/BlockNoteExtension.js";
+import { YCursor } from "./YCursorPlugin.js";
+import { YSync } from "./YSync.js";
+import { YUndo } from "./YUndo.js";
+
+/**
+ * To find a fragment in another ydoc, we need to search for it.
+ */
+function findTypeInOtherYdoc<T extends Y.AbstractType<any>>(
+  ytype: T,
+  otherYdoc: Y.Doc,
+): T {
+  const ydoc = ytype.doc!;
+  if (ytype._item === null) {
+    /**
+     * If is a root type, we need to find the root key in the original ydoc
+     * and use it to get the type in the other ydoc.
+     */
+    const rootKey = Array.from(ydoc.share.keys()).find(
+      (key) => ydoc.share.get(key) === ytype,
+    );
+    if (rootKey == null) {
+      throw new Error("type does not exist in other ydoc");
+    }
+    return otherYdoc.get(rootKey, ytype.constructor as new () => T) as T;
+  } else {
+    /**
+     * If it is a sub type, we use the item id to find the history type.
+     */
+    const ytypeItem = ytype._item;
+    const otherStructs = otherYdoc.store.clients.get(ytypeItem.id.client) ?? [];
+    const itemIndex = Y.findIndexSS(otherStructs, ytypeItem.id.clock);
+    const otherItem = otherStructs[itemIndex] as Y.Item;
+    const otherContent = otherItem.content as Y.ContentType;
+    return otherContent.type as T;
+  }
+}
+
+export const ForkYDoc = createExtension((editor, options) => {
+  if (!options.collaboration) {
+    return;
+  }
+
+  let forkedState:
+    | {
+        originalFragment: Y.XmlFragment;
+        undoStack: Y.UndoManager["undoStack"];
+        forkedFragment: Y.XmlFragment;
+      }
+    | undefined = undefined;
+
+  const store = createStore({ isForked: false });
+
+  return {
+    key: "yForkDoc",
+    store,
+    /**
+     * Fork the Y.js document from syncing to the remote,
+     * allowing modifications to the document without affecting the remote.
+     * These changes can later be rolled back or applied to the remote.
+     */
+    fork() {
+      if (forkedState) {
+        return;
+      }
+
+      const originalFragment = options.collaboration?.fragment;
+
+      if (!originalFragment) {
+        throw new Error("No fragment to fork from");
+      }
+
+      const doc = new Y.Doc();
+      // Copy the original document to a new Yjs document
+      Y.applyUpdate(doc, Y.encodeStateAsUpdate(originalFragment.doc!));
+
+      // Find the forked fragment in the new Yjs document
+      const forkedFragment = findTypeInOtherYdoc(originalFragment, doc);
+
+      forkedState = {
+        undoStack: yUndoPluginKey.getState(editor.prosemirrorState)!.undoManager
+          .undoStack,
+        originalFragment,
+        forkedFragment,
+      };
+
+      // Need to reset all the yjs plugins
+      editor.removeExtension([YUndo, YCursor, YSync]);
+      const newOptions = {
+        ...options,
+        collaboration: {
+          ...options.collaboration!,
+          fragment: forkedFragment,
+        },
+      };
+      // Register them again, based on the new forked fragment
+      editor.registerExtension([
+        YSync(editor, newOptions),
+        // No need to register the cursor plugin again, it's a local fork
+        YUndo(editor, newOptions),
+      ]);
+
+      // Tell the store that the editor is now forked
+      store.setState({ isForked: true });
+    },
+
+    /**
+     * Resume syncing the Y.js document to the remote
+     * If `keepChanges` is true, any changes that have been made to the forked document will be applied to the original document.
+     * Otherwise, the original document will be restored and the changes will be discarded.
+     */
+    merge({ keepChanges }: { keepChanges: boolean }) {
+      if (!forkedState) {
+        return;
+      }
+      // Remove the forked fragment's plugins
+      editor.removeExtension(["ySync", "yCursor", "yUndo"]);
+
+      const { originalFragment, forkedFragment, undoStack } = forkedState;
+      // Register the plugins again, based on the original fragment (which is still in the original options)
+      editor.registerExtension([YSync, YCursor, YUndo]);
+
+      // Reset the undo stack to the original undo stack
+      yUndoPluginKey.getState(editor.prosemirrorState)!.undoManager.undoStack =
+        undoStack;
+
+      if (keepChanges) {
+        // Apply any changes that have been made to the fork, onto the original doc
+        const update = Y.encodeStateAsUpdate(
+          forkedFragment.doc!,
+          Y.encodeStateVector(originalFragment.doc!),
+        );
+        // Applying this change will add to the undo stack, allowing it to be undone normally
+        Y.applyUpdate(originalFragment.doc!, update, editor);
+      }
+      // Reset the forked state
+      forkedState = undefined;
+      // Tell the store that the editor is no longer forked
+      store.setState({ isForked: false });
+    },
+  } as const;
+});
diff --git a/packages/core/src/extensions/Collaboration/ForkYDocPlugin.ts b/packages/core/src/extensions/Collaboration/ForkYDocPlugin.ts
deleted file mode 100644
index c6361557e7..0000000000
--- a/packages/core/src/extensions/Collaboration/ForkYDocPlugin.ts
+++ /dev/null
@@ -1,192 +0,0 @@
-import * as Y from "yjs";
-
-import {
-  yCursorPluginKey,
-  ySyncPluginKey,
-  yUndoPluginKey,
-} from "y-prosemirror";
-import { CursorPlugin } from "./CursorPlugin.js";
-import { SyncPlugin } from "./SyncPlugin.js";
-import { UndoPlugin } from "./UndoPlugin.js";
-
-import {
-  BlockNoteEditor,
-  BlockNoteEditorOptions,
-} from "../../editor/BlockNoteEditor.js";
-import { BlockNoteExtension } from "../../editor/BlockNoteExtension.js";
-
-export class ForkYDocPlugin extends BlockNoteExtension<{
-  forked: boolean;
-}> {
-  public static key() {
-    return "ForkYDocPlugin";
-  }
-
-  private editor: BlockNoteEditor<any, any, any>;
-  private collaboration: BlockNoteEditorOptions<any, any, any>["collaboration"];
-
-  constructor({
-    editor,
-    collaboration,
-  }: {
-    editor: BlockNoteEditor<any, any, any>;
-    collaboration: BlockNoteEditorOptions<any, any, any>["collaboration"];
-  }) {
-    super(editor);
-    this.editor = editor;
-    this.collaboration = collaboration;
-  }
-
-  /**
-   * To find a fragment in another ydoc, we need to search for it.
-   */
-  private findTypeInOtherYdoc<T extends Y.AbstractType<any>>(
-    ytype: T,
-    otherYdoc: Y.Doc,
-  ): T {
-    const ydoc = ytype.doc!;
-    if (ytype._item === null) {
-      /**
-       * If is a root type, we need to find the root key in the original ydoc
-       * and use it to get the type in the other ydoc.
-       */
-      const rootKey = Array.from(ydoc.share.keys()).find(
-        (key) => ydoc.share.get(key) === ytype,
-      );
-      if (rootKey == null) {
-        throw new Error("type does not exist in other ydoc");
-      }
-      return otherYdoc.get(rootKey, ytype.constructor as new () => T) as T;
-    } else {
-      /**
-       * If it is a sub type, we use the item id to find the history type.
-       */
-      const ytypeItem = ytype._item;
-      const otherStructs =
-        otherYdoc.store.clients.get(ytypeItem.id.client) ?? [];
-      const itemIndex = Y.findIndexSS(otherStructs, ytypeItem.id.clock);
-      const otherItem = otherStructs[itemIndex] as Y.Item;
-      const otherContent = otherItem.content as Y.ContentType;
-      return otherContent.type as T;
-    }
-  }
-
-  /**
-   * Whether the editor is editing a forked document,
-   * preserving a reference to the original document and the forked document.
-   */
-  public get isForkedFromRemote() {
-    return this.forkedState !== undefined;
-  }
-
-  /**
-   * Stores whether the editor is editing a forked document,
-   * preserving a reference to the original document and the forked document.
-   */
-  private forkedState:
-    | {
-        originalFragment: Y.XmlFragment;
-        undoStack: Y.UndoManager["undoStack"];
-        forkedFragment: Y.XmlFragment;
-      }
-    | undefined;
-
-  /**
-   * Fork the Y.js document from syncing to the remote,
-   * allowing modifications to the document without affecting the remote.
-   * These changes can later be rolled back or applied to the remote.
-   */
-  public fork() {
-    if (this.isForkedFromRemote) {
-      return;
-    }
-
-    const originalFragment = this.collaboration?.fragment;
-
-    if (!originalFragment) {
-      throw new Error("No fragment to fork from");
-    }
-
-    const doc = new Y.Doc();
-    // Copy the original document to a new Yjs document
-    Y.applyUpdate(doc, Y.encodeStateAsUpdate(originalFragment.doc!));
-
-    // Find the forked fragment in the new Yjs document
-    const forkedFragment = this.findTypeInOtherYdoc(originalFragment, doc);
-
-    this.forkedState = {
-      undoStack: yUndoPluginKey.getState(this.editor.prosemirrorState)!
-        .undoManager.undoStack,
-      originalFragment,
-      forkedFragment,
-    };
-
-    // Need to reset all the yjs plugins
-    this.editor._tiptapEditor.unregisterPlugin([
-      yCursorPluginKey,
-      yUndoPluginKey,
-      ySyncPluginKey,
-    ]);
-    // Register them again, based on the new forked fragment
-    this.editor._tiptapEditor.registerPlugin(
-      new SyncPlugin(forkedFragment).plugins[0],
-    );
-    this.editor._tiptapEditor.registerPlugin(
-      new UndoPlugin({ editor: this.editor }).plugins[0],
-    );
-    // No need to register the cursor plugin again, it's a local fork
-    this.emit("forked", true);
-  }
-
-  /**
-   * Resume syncing the Y.js document to the remote
-   * If `keepChanges` is true, any changes that have been made to the forked document will be applied to the original document.
-   * Otherwise, the original document will be restored and the changes will be discarded.
-   */
-  public merge({ keepChanges }: { keepChanges: boolean }) {
-    if (!this.forkedState) {
-      return;
-    }
-    // Remove the forked fragment's plugins
-    this.editor._tiptapEditor.unregisterPlugin(ySyncPluginKey);
-    this.editor._tiptapEditor.unregisterPlugin(yUndoPluginKey);
-
-    const { originalFragment, forkedFragment, undoStack } = this.forkedState;
-    this.editor.extensions["ySyncPlugin"] = new SyncPlugin(originalFragment);
-    this.editor.extensions["yCursorPlugin"] = new CursorPlugin(
-      this.collaboration!,
-    );
-    this.editor.extensions["yUndoPlugin"] = new UndoPlugin({
-      editor: this.editor,
-    });
-
-    // Register the plugins again, based on the original fragment
-    this.editor._tiptapEditor.registerPlugin(
-      this.editor.extensions["ySyncPlugin"].plugins[0],
-    );
-    this.editor._tiptapEditor.registerPlugin(
-      this.editor.extensions["yCursorPlugin"].plugins[0],
-    );
-    this.editor._tiptapEditor.registerPlugin(
-      this.editor.extensions["yUndoPlugin"].plugins[0],
-    );
-
-    // Reset the undo stack to the original undo stack
-    yUndoPluginKey.getState(
-      this.editor.prosemirrorState,
-    )!.undoManager.undoStack = undoStack;
-
-    if (keepChanges) {
-      // Apply any changes that have been made to the fork, onto the original doc
-      const update = Y.encodeStateAsUpdate(
-        forkedFragment.doc!,
-        Y.encodeStateVector(originalFragment.doc!),
-      );
-      // Applying this change will add to the undo stack, allowing it to be undone normally
-      Y.applyUpdate(originalFragment.doc!, update, this.editor);
-    }
-    // Reset the forked state
-    this.forkedState = undefined;
-    this.emit("forked", false);
-  }
-}
diff --git a/packages/core/src/extensions/Collaboration/SyncPlugin.ts b/packages/core/src/extensions/Collaboration/SyncPlugin.ts
deleted file mode 100644
index 15ec5751c8..0000000000
--- a/packages/core/src/extensions/Collaboration/SyncPlugin.ts
+++ /dev/null
@@ -1,18 +0,0 @@
-import { ySyncPlugin } from "y-prosemirror";
-import type * as Y from "yjs";
-import { BlockNoteExtension } from "../../editor/BlockNoteExtension.js";
-
-export class SyncPlugin extends BlockNoteExtension {
-  public static key() {
-    return "ySyncPlugin";
-  }
-
-  constructor(fragment: Y.XmlFragment) {
-    super();
-    this.addProsemirrorPlugin(ySyncPlugin(fragment));
-  }
-
-  public get priority() {
-    return 1001;
-  }
-}
diff --git a/packages/core/src/extensions/Collaboration/UndoPlugin.ts b/packages/core/src/extensions/Collaboration/UndoPlugin.ts
deleted file mode 100644
index e0e697de08..0000000000
--- a/packages/core/src/extensions/Collaboration/UndoPlugin.ts
+++ /dev/null
@@ -1,18 +0,0 @@
-import { yUndoPlugin } from "y-prosemirror";
-import { BlockNoteEditor } from "../../editor/BlockNoteEditor.js";
-import { BlockNoteExtension } from "../../editor/BlockNoteExtension.js";
-
-export class UndoPlugin extends BlockNoteExtension {
-  public static key() {
-    return "yUndoPlugin";
-  }
-
-  constructor({ editor }: { editor: BlockNoteEditor<any, any, any> }) {
-    super();
-    this.addProsemirrorPlugin(yUndoPlugin({ trackedOrigins: [editor] }));
-  }
-
-  public get priority() {
-    return 1000;
-  }
-}
diff --git a/packages/core/src/extensions/Collaboration/YCursorPlugin.ts b/packages/core/src/extensions/Collaboration/YCursorPlugin.ts
new file mode 100644
index 0000000000..6bbd3d0e5c
--- /dev/null
+++ b/packages/core/src/extensions/Collaboration/YCursorPlugin.ts
@@ -0,0 +1,180 @@
+import { defaultSelectionBuilder, yCursorPlugin } from "y-prosemirror";
+import { createExtension } from "../../editor/BlockNoteExtension.js";
+
+export type CollaborationUser = {
+  name: string;
+  color: string;
+  [key: string]: string;
+};
+
+/**
+ * Determine whether the foreground color should be white or black based on a provided background color
+ * Inspired by: https://stackoverflow.com/a/3943023
+ */
+function isDarkColor(bgColor: string): boolean {
+  const color = bgColor.charAt(0) === "#" ? bgColor.substring(1, 7) : bgColor;
+  const r = parseInt(color.substring(0, 2), 16); // hexToR
+  const g = parseInt(color.substring(2, 4), 16); // hexToG
+  const b = parseInt(color.substring(4, 6), 16); // hexToB
+  const uicolors = [r / 255, g / 255, b / 255];
+  const c = uicolors.map((col) => {
+    if (col <= 0.03928) {
+      return col / 12.92;
+    }
+    return Math.pow((col + 0.055) / 1.055, 2.4);
+  });
+  const L = 0.2126 * c[0] + 0.7152 * c[1] + 0.0722 * c[2];
+  return L <= 0.179;
+}
+
+function defaultCursorRender(user: CollaborationUser) {
+  const cursorElement = document.createElement("span");
+
+  cursorElement.classList.add("bn-collaboration-cursor__base");
+
+  const caretElement = document.createElement("span");
+  caretElement.setAttribute("contentedEditable", "false");
+  caretElement.classList.add("bn-collaboration-cursor__caret");
+  caretElement.setAttribute(
+    "style",
+    `background-color: ${user.color}; color: ${
+      isDarkColor(user.color) ? "white" : "black"
+    }`,
+  );
+
+  const labelElement = document.createElement("span");
+
+  labelElement.classList.add("bn-collaboration-cursor__label");
+  labelElement.setAttribute(
+    "style",
+    `background-color: ${user.color}; color: ${
+      isDarkColor(user.color) ? "white" : "black"
+    }`,
+  );
+  labelElement.insertBefore(document.createTextNode(user.name), null);
+
+  caretElement.insertBefore(labelElement, null);
+
+  cursorElement.insertBefore(document.createTextNode("\u2060"), null); // Non-breaking space
+  cursorElement.insertBefore(caretElement, null);
+  cursorElement.insertBefore(document.createTextNode("\u2060"), null); // Non-breaking space
+
+  return cursorElement;
+}
+
+export const YCursor = createExtension((_editor, options) => {
+  const collaboration = options?.collaboration;
+  if (!collaboration) {
+    return;
+  }
+
+  const recentlyUpdatedCursors = new Map();
+
+  if (
+    "awareness" in collaboration.provider &&
+    typeof collaboration.provider.awareness === "object"
+  ) {
+    if (
+      "setLocalStateField" in collaboration.provider.awareness &&
+      typeof collaboration.provider.awareness.setLocalStateField === "function"
+    ) {
+      collaboration.provider.awareness.setLocalStateField(
+        "user",
+        collaboration.user,
+      );
+    }
+    if (
+      "on" in collaboration.provider.awareness &&
+      typeof collaboration.provider.awareness.on === "function"
+    ) {
+      if (collaboration.showCursorLabels !== "always") {
+        collaboration.provider.awareness.on(
+          "change",
+          ({
+            updated,
+          }: {
+            added: Array<number>;
+            updated: Array<number>;
+            removed: Array<number>;
+          }) => {
+            for (const clientID of updated) {
+              const cursor = recentlyUpdatedCursors.get(clientID);
+
+              if (cursor) {
+                cursor.element.setAttribute("data-active", "");
+
+                if (cursor.hideTimeout) {
+                  clearTimeout(cursor.hideTimeout);
+                }
+
+                recentlyUpdatedCursors.set(clientID, {
+                  element: cursor.element,
+                  hideTimeout: setTimeout(() => {
+                    cursor.element.removeAttribute("data-active");
+                  }, 2000),
+                });
+              }
+            }
+          },
+        );
+      }
+    }
+  }
+
+  return {
+    key: "yCursor",
+    prosemirrorPlugins: [
+      yCursorPlugin(collaboration.provider.awareness, {
+        selectionBuilder: defaultSelectionBuilder,
+        cursorBuilder(user: CollaborationUser, clientID: number) {
+          let cursorData = recentlyUpdatedCursors.get(clientID);
+
+          if (!cursorData) {
+            const cursorElement = (
+              collaboration.renderCursor ?? defaultCursorRender
+            )(user);
+
+            if (collaboration.showCursorLabels !== "always") {
+              cursorElement.addEventListener("mouseenter", () => {
+                const cursor = recentlyUpdatedCursors.get(clientID)!;
+                cursor.element.setAttribute("data-active", "");
+
+                if (cursor.hideTimeout) {
+                  clearTimeout(cursor.hideTimeout);
+                  recentlyUpdatedCursors.set(clientID, {
+                    element: cursor.element,
+                    hideTimeout: undefined,
+                  });
+                }
+              });
+
+              cursorElement.addEventListener("mouseleave", () => {
+                const cursor = recentlyUpdatedCursors.get(clientID)!;
+
+                recentlyUpdatedCursors.set(clientID, {
+                  element: cursor.element,
+                  hideTimeout: setTimeout(() => {
+                    cursor.element.removeAttribute("data-active");
+                  }, 2000),
+                });
+              });
+            }
+
+            cursorData = {
+              element: cursorElement,
+              hideTimeout: undefined,
+            };
+
+            recentlyUpdatedCursors.set(clientID, cursorData);
+          }
+
+          return cursorData.element;
+        },
+      }),
+    ],
+    dependsOn: ["ySync"],
+    updateUser(user: { name: string; color: string; [key: string]: string }) {
+      collaboration.provider.awareness.setLocalStateField("user", user);
+    },
+  } as const;
+});
diff --git a/packages/core/src/extensions/Collaboration/YSync.ts b/packages/core/src/extensions/Collaboration/YSync.ts
new file mode 100644
index 0000000000..14c5e3f023
--- /dev/null
+++ b/packages/core/src/extensions/Collaboration/YSync.ts
@@ -0,0 +1,15 @@
+import { ySyncPlugin } from "y-prosemirror";
+import { createExtension } from "../../editor/BlockNoteExtension.js";
+
+export const YSync = createExtension((_editor, options) => {
+  const fragment = options?.collaboration?.fragment;
+  if (!fragment) {
+    return;
+  }
+
+  return {
+    key: "ySync",
+    prosemirrorPlugins: [ySyncPlugin(fragment)],
+    runsBefore: ["default"],
+  } as const;
+});
diff --git a/packages/core/src/extensions/Collaboration/YUndo.ts b/packages/core/src/extensions/Collaboration/YUndo.ts
new file mode 100644
index 0000000000..9a797bcfb0
--- /dev/null
+++ b/packages/core/src/extensions/Collaboration/YUndo.ts
@@ -0,0 +1,16 @@
+import { redoCommand, undoCommand, yUndoPlugin } from "y-prosemirror";
+import { createExtension } from "../../editor/BlockNoteExtension.js";
+
+export const YUndo = createExtension((editor, options) => {
+  if (!options.collaboration) {
+    return;
+  }
+
+  return {
+    key: "yUndo",
+    prosemirrorPlugins: [yUndoPlugin({ trackedOrigins: [editor] })],
+    dependsOn: ["yCursor", "ySync"],
+    undoCommand: undoCommand,
+    redoCommand: redoCommand,
+  } as const;
+});
diff --git a/packages/core/src/extensions/Collaboration/schemaMigration/SchemaMigrationPlugin.ts b/packages/core/src/extensions/Collaboration/schemaMigration/SchemaMigration.ts
similarity index 75%
rename from packages/core/src/extensions/Collaboration/schemaMigration/SchemaMigrationPlugin.ts
rename to packages/core/src/extensions/Collaboration/schemaMigration/SchemaMigration.ts
index fc2962dd6b..21393556a5 100644
--- a/packages/core/src/extensions/Collaboration/schemaMigration/SchemaMigrationPlugin.ts
+++ b/packages/core/src/extensions/Collaboration/schemaMigration/SchemaMigration.ts
@@ -1,7 +1,7 @@
 import { Plugin, PluginKey } from "@tiptap/pm/state";
 import * as Y from "yjs";
 
-import { BlockNoteExtension } from "../../../editor/BlockNoteExtension.js";
+import { createExtension } from "../../../editor/BlockNoteExtension.js";
 import migrationRules from "./migrationRules/index.js";
 
 // This plugin allows us to update collaboration YDocs whenever BlockNote's
@@ -10,22 +10,25 @@ import migrationRules from "./migrationRules/index.js";
 // case things are found in the fragment that don't adhere to the editor schema
 // and need to be fixed. These fixes are defined as `MigrationRule`s within the
 // `migrationRules` directory.
-export class SchemaMigrationPlugin extends BlockNoteExtension {
-  private migrationDone = false;
+export const SchemaMigration = createExtension((_editor, options) => {
+  const fragment = (options as any)?.collaboration?.fragment as
+    | Y.XmlFragment
+    | undefined;
 
-  public static key() {
-    return "schemaMigrationPlugin";
+  if (!fragment) {
+    return;
   }
 
-  constructor(fragment: Y.XmlFragment) {
-    const pluginKey = new PluginKey(SchemaMigrationPlugin.key());
+  let migrationDone = false;
+  const pluginKey = new PluginKey("schemaMigration");
 
-    super();
-    this.addProsemirrorPlugin(
+  return {
+    key: "schemaMigration",
+    prosemirrorPlugins: [
       new Plugin({
         key: pluginKey,
         appendTransaction: (transactions, _oldState, newState) => {
-          if (this.migrationDone) {
+          if (migrationDone) {
             return undefined;
           }
 
@@ -45,7 +48,7 @@ export class SchemaMigrationPlugin extends BlockNoteExtension {
             migrationRule(fragment, tr);
           }
 
-          this.migrationDone = true;
+          migrationDone = true;
 
           if (!tr.docChanged) {
             return undefined;
@@ -54,6 +57,6 @@ export class SchemaMigrationPlugin extends BlockNoteExtension {
           return tr;
         },
       }),
-    );
-  }
-}
+    ],
+  } as const;
+});
diff --git a/packages/core/src/extensions/Comments/Comments.ts b/packages/core/src/extensions/Comments/Comments.ts
new file mode 100644
index 0000000000..697fe48f39
--- /dev/null
+++ b/packages/core/src/extensions/Comments/Comments.ts
@@ -0,0 +1,327 @@
+import { Node } from "prosemirror-model";
+import { Plugin, PluginKey } from "prosemirror-state";
+import { Decoration, DecorationSet } from "prosemirror-view";
+import { getRelativeSelection, ySyncPluginKey } from "y-prosemirror";
+import type { CommentBody, ThreadData, User } from "../../comments/index.js";
+import type { BlockNoteEditor } from "../../editor/BlockNoteEditor.js";
+import {
+  createExtension,
+  createStore,
+} from "../../editor/BlockNoteExtension.js";
+import { UserStore } from "./userstore/UserStore.js";
+import { CommentMark } from "./CommentMark.js";
+
+const PLUGIN_KEY = new PluginKey(`blocknote-comments`);
+const SET_SELECTED_THREAD_ID = "SET_SELECTED_THREAD_ID";
+
+type CommentsPluginState = {
+  /**
+   * Decorations to be rendered, specifically to indicate the selected thread
+   */
+  decorations: DecorationSet;
+};
+
+/**
+ * Calculate the thread positions from the current document state
+ */
+function getUpdatedThreadPositions(doc: Node, markType: string) {
+  const threadPositions = new Map<string, { from: number; to: number }>();
+
+  // find all thread marks and store their position + create decoration for selected thread
+  doc.descendants((node, pos) => {
+    node.marks.forEach((mark) => {
+      if (mark.type.name === markType) {
+        const thisThreadId = (mark.attrs as { threadId: string | undefined })
+          .threadId;
+        if (!thisThreadId) {
+          return;
+        }
+        const from = pos;
+        const to = from + node.nodeSize;
+
+        // FloatingThreads component uses "to" as the position, so always store the largest "to" found
+        // AnchoredThreads component uses "from" as the position, so always store the smallest "from" found
+        const currentPosition = threadPositions.get(thisThreadId) ?? {
+          from: Infinity,
+          to: 0,
+        };
+        threadPositions.set(thisThreadId, {
+          from: Math.min(from, currentPosition.from),
+          to: Math.max(to, currentPosition.to),
+        });
+      }
+    });
+  });
+  return threadPositions;
+}
+
+export const Comments = createExtension(
+  (editor: BlockNoteEditor<any, any, any>, options) => {
+    const commentsOptions = options?.comments;
+    if (!commentsOptions) {
+      return undefined;
+    }
+
+    const markType = CommentMark.name;
+    const threadStore = commentsOptions.threadStore;
+    const resolveUsers = options.resolveUsers;
+    const commentEditorSchema = commentsOptions.schema;
+
+    if (!resolveUsers) {
+      throw new Error("resolveUsers is required for comments");
+    }
+
+    const userStore = new UserStore<User>(resolveUsers);
+    let pendingComment = false;
+    let selectedThreadId: string | undefined;
+    let threadPositions: Map<string, { from: number; to: number }> = new Map();
+    const store = createStore({
+      pendingComment,
+      selectedThreadId,
+      threadPositions,
+    });
+    const subscribers = new Set<
+      (state: {
+        pendingComment: boolean;
+        selectedThreadId: string | undefined;
+        threadPositions: Map<string, { from: number; to: number }>;
+      }) => void
+    >();
+
+    const emitStateUpdate = () => {
+      const snapshot = {
+        pendingComment,
+        selectedThreadId,
+        threadPositions,
+      };
+      store.setState(snapshot);
+      subscribers.forEach((cb) => cb(snapshot));
+    };
+
+    const updateMarksFromThreads = (threads: Map<string, ThreadData>) => {
+      editor.transact((tr) => {
+        tr.doc.descendants((node, pos) => {
+          node.marks.forEach((mark) => {
+            if (mark.type.name === markType) {
+              const markTypeInstance = mark.type;
+              const markThreadId = mark.attrs.threadId;
+              const thread = threads.get(markThreadId);
+              const isOrphan = !!(
+                !thread ||
+                thread.resolved ||
+                thread.deletedAt
+              );
+
+              if (isOrphan !== mark.attrs.orphan) {
+                const trimmedFrom = Math.max(pos, 0);
+                const trimmedTo = Math.min(
+                  pos + node.nodeSize,
+                  tr.doc.content.size - 1,
+                  tr.doc.content.size - 1,
+                );
+                tr.removeMark(trimmedFrom, trimmedTo, mark);
+                tr.addMark(
+                  trimmedFrom,
+                  trimmedTo,
+                  markTypeInstance.create({
+                    ...mark.attrs,
+                    orphan: isOrphan,
+                  }),
+                );
+
+                if (isOrphan && selectedThreadId === markThreadId) {
+                  // unselect
+                  selectedThreadId = undefined;
+                  emitStateUpdate();
+                }
+              }
+            }
+          });
+        });
+      });
+    };
+
+    return {
+      key: "comments",
+      store,
+      prosemirrorPlugins: [
+        new Plugin<CommentsPluginState>({
+          key: PLUGIN_KEY,
+          state: {
+            init() {
+              return {
+                decorations: DecorationSet.empty,
+              };
+            },
+            apply(tr, state) {
+              const action = tr.getMeta(PLUGIN_KEY);
+
+              if (!tr.docChanged && !action) {
+                return state;
+              }
+
+              // only update threadPositions if the doc changed
+              const newThreadPositions = tr.docChanged
+                ? getUpdatedThreadPositions(tr.doc, markType)
+                : threadPositions;
+
+              if (newThreadPositions.size > 0 || threadPositions.size > 0) {
+                // small optimization; don't emit event if threadPositions before / after were both empty
+                threadPositions = newThreadPositions;
+                emitStateUpdate();
+              }
+
+              // update decorations if doc or selected thread changed
+              const decorations = [] as any[];
+
+              if (selectedThreadId) {
+                const selectedThreadPosition =
+                  newThreadPositions.get(selectedThreadId);
+
+                if (selectedThreadPosition) {
+                  decorations.push(
+                    Decoration.inline(
+                      selectedThreadPosition.from,
+                      selectedThreadPosition.to,
+                      {
+                        class: "bn-thread-mark-selected",
+                      },
+                    ),
+                  );
+                }
+              }
+
+              return {
+                decorations: DecorationSet.create(tr.doc, decorations),
+              };
+            },
+          },
+          props: {
+            decorations(state) {
+              return (
+                PLUGIN_KEY.getState(state)?.decorations ?? DecorationSet.empty
+              );
+            },
+            handleClick: (view, pos, event) => {
+              if (event.button !== 0) {
+                return;
+              }
+
+              const node = view.state.doc.nodeAt(pos);
+
+              if (!node) {
+                // unselect
+                selectedThreadId = undefined;
+                emitStateUpdate();
+                return;
+              }
+
+              const commentMark = node.marks.find(
+                (mark) =>
+                  mark.type.name === markType && mark.attrs.orphan !== true,
+              );
+
+              const threadId = commentMark?.attrs.threadId as
+                | string
+                | undefined;
+              if (threadId !== selectedThreadId) {
+                selectedThreadId = threadId;
+                emitStateUpdate();
+              }
+            },
+          },
+        }),
+      ],
+      threadStore: threadStore,
+      init() {
+        const unsubscribe = threadStore.subscribe(updateMarksFromThreads);
+        editor.onCreate(() => {
+          updateMarksFromThreads(threadStore.getThreads());
+          editor.onSelectionChange(() => {
+            if (pendingComment) {
+              pendingComment = false;
+              emitStateUpdate();
+            }
+          });
+        });
+        return () => unsubscribe();
+      },
+      onUpdate(
+        callback: (state: {
+          pendingComment: boolean;
+          selectedThreadId: string | undefined;
+          threadPositions: Map<string, { from: number; to: number }>;
+        }) => void,
+      ) {
+        subscribers.add(callback);
+        return () => subscribers.delete(callback);
+      },
+      selectThread(threadId: string | undefined, scrollToThread = true) {
+        if (selectedThreadId === threadId) {
+          return;
+        }
+        selectedThreadId = threadId;
+        emitStateUpdate();
+        editor.transact((tr) =>
+          tr.setMeta(PLUGIN_KEY, {
+            name: SET_SELECTED_THREAD_ID,
+          }),
+        );
+
+        if (threadId && scrollToThread) {
+          const selectedThreadPosition = threadPositions.get(threadId);
+          if (!selectedThreadPosition) {
+            return;
+          }
+          (
+            editor.prosemirrorView?.domAtPos(selectedThreadPosition.from)
+              .node as Element | undefined
+          )?.scrollIntoView({
+            behavior: "smooth",
+            block: "center",
+          });
+        }
+      },
+      startPendingComment() {
+        pendingComment = true;
+        emitStateUpdate();
+      },
+      stopPendingComment() {
+        pendingComment = false;
+        emitStateUpdate();
+      },
+      async createThread(options: {
+        initialComment: { body: CommentBody; metadata?: any };
+        metadata?: any;
+      }) {
+        const thread = await threadStore.createThread(options);
+        if (threadStore.addThreadToDocument) {
+          const view = editor.prosemirrorView!;
+          const pmSelection = view.state.selection;
+          const ystate = ySyncPluginKey.getState(view.state);
+          const selection = {
+            prosemirror: {
+              head: pmSelection.head,
+              anchor: pmSelection.anchor,
+            },
+            yjs: ystate
+              ? getRelativeSelection(ystate.binding, view.state)
+              : undefined,
+          };
+          await threadStore.addThreadToDocument({
+            threadId: thread.id,
+            selection,
+          });
+        } else {
+          (editor as any)._tiptapEditor.commands.setMark(markType, {
+            orphan: false,
+            threadId: thread.id,
+          });
+        }
+      },
+      userStore,
+      commentEditorSchema,
+      tiptapExtensions: [CommentMark],
+    } as const;
+  },
+);
diff --git a/packages/core/src/extensions/Comments/CommentsPlugin.ts b/packages/core/src/extensions/Comments/CommentsPlugin.ts
deleted file mode 100644
index 23eea69e92..0000000000
--- a/packages/core/src/extensions/Comments/CommentsPlugin.ts
+++ /dev/null
@@ -1,392 +0,0 @@
-import { Node } from "prosemirror-model";
-import { Plugin, PluginKey } from "prosemirror-state";
-import { Decoration, DecorationSet } from "prosemirror-view";
-import { getRelativeSelection, ySyncPluginKey } from "y-prosemirror";
-import type {
-  CommentBody,
-  ThreadData,
-  ThreadStore,
-  User,
-} from "../../comments/index.js";
-import { BlockNoteEditor } from "../../editor/BlockNoteEditor.js";
-import { BlockNoteExtension } from "../../editor/BlockNoteExtension.js";
-import { CustomBlockNoteSchema } from "../../schema/schema.js";
-import { UserStore } from "./userstore/UserStore.js";
-import { getMarkRange } from "@tiptap/core";
-
-const PLUGIN_KEY = new PluginKey(`blocknote-comments`);
-const SET_SELECTED_THREAD_ID = "SET_SELECTED_THREAD_ID";
-
-type CommentsPluginState = {
-  /**
-   * Decorations to be rendered, specifically to indicate the selected thread
-   */
-  decorations: DecorationSet;
-};
-
-/**
- * Calculate the thread positions from the current document state
- */
-function getUpdatedThreadPositions(doc: Node, markType: string) {
-  const threadPositions = new Map<string, { from: number; to: number }>();
-
-  // find all thread marks and store their position + create decoration for selected thread
-  doc.descendants((node, pos) => {
-    node.marks.forEach((mark) => {
-      if (mark.type.name === markType) {
-        const thisThreadId = (mark.attrs as { threadId: string | undefined })
-          .threadId;
-        if (!thisThreadId) {
-          return;
-        }
-        const from = pos;
-        const to = from + node.nodeSize;
-
-        // FloatingThreads component uses "to" as the position, so always store the largest "to" found
-        // AnchoredThreads component uses "from" as the position, so always store the smallest "from" found
-        const currentPosition = threadPositions.get(thisThreadId) ?? {
-          from: Infinity,
-          to: 0,
-        };
-        threadPositions.set(thisThreadId, {
-          from: Math.min(from, currentPosition.from),
-          to: Math.max(to, currentPosition.to),
-        });
-      }
-    });
-  });
-  return threadPositions;
-}
-
-export class CommentsPlugin extends BlockNoteExtension {
-  public static key() {
-    return "comments";
-  }
-
-  public readonly userStore: UserStore<User>;
-
-  /**
-   * Whether a comment is currently being composed
-   */
-  private pendingComment = false;
-
-  /**
-   * The currently selected thread id
-   */
-  private selectedThreadId: string | undefined;
-
-  /**
-   * Store the positions of all threads in the document.
-   * this can be used later to implement a floating sidebar
-   */
-  private threadPositions: Map<string, { from: number; to: number }> =
-    new Map();
-
-  private emitStateUpdate() {
-    this.emit("update", {
-      selectedThreadId: this.selectedThreadId,
-      pendingComment: this.pendingComment,
-      threadPositions: this.threadPositions,
-    });
-  }
-
-  /**
-   * when a thread is resolved or deleted, we need to update the marks to reflect the new state
-   */
-  private updateMarksFromThreads = (threads: Map<string, ThreadData>) => {
-    this.editor.transact((tr) => {
-      tr.doc.descendants((node, pos) => {
-        node.marks.forEach((mark) => {
-          if (mark.type.name === this.markType) {
-            const markType = mark.type;
-            const markThreadId = mark.attrs.threadId;
-            const thread = threads.get(markThreadId);
-            const isOrphan = !!(!thread || thread.resolved || thread.deletedAt);
-
-            if (isOrphan !== mark.attrs.orphan) {
-              const trimmedFrom = Math.max(pos, 0);
-              const trimmedTo = Math.min(
-                pos + node.nodeSize,
-                tr.doc.content.size - 1,
-                tr.doc.content.size - 1,
-              );
-              tr.removeMark(trimmedFrom, trimmedTo, mark);
-              tr.addMark(
-                trimmedFrom,
-                trimmedTo,
-                markType.create({
-                  ...mark.attrs,
-                  orphan: isOrphan,
-                }),
-              );
-
-              if (isOrphan && this.selectedThreadId === markThreadId) {
-                // unselect
-                this.selectedThreadId = undefined;
-                this.emitStateUpdate();
-              }
-            }
-          }
-        });
-      });
-    });
-  };
-
-  constructor(
-    private readonly editor: BlockNoteEditor<any, any, any>,
-    public readonly threadStore: ThreadStore,
-    private readonly markType: string,
-    public readonly resolveUsers:
-      | undefined
-      | ((userIds: string[]) => Promise<User[]>),
-    public readonly commentEditorSchema?: CustomBlockNoteSchema<any, any, any>,
-  ) {
-    super();
-
-    if (!resolveUsers) {
-      throw new Error("resolveUsers is required for comments");
-    }
-    this.userStore = new UserStore<User>(resolveUsers);
-
-    // Note: Plugins are currently not destroyed when the editor is destroyed.
-    // We should unsubscribe from the threadStore when the editor is destroyed.
-    this.threadStore.subscribe(this.updateMarksFromThreads);
-
-    editor.onCreate(() => {
-      // Need to wait for TipTap editor state to be initialized
-      this.updateMarksFromThreads(this.threadStore.getThreads());
-      editor.onSelectionChange(() => {
-        if (this.pendingComment) {
-          this.pendingComment = false;
-          this.emitStateUpdate();
-        }
-      });
-    });
-
-    // eslint-disable-next-line @typescript-eslint/no-this-alias
-    const self = this;
-
-    this.addProsemirrorPlugin(
-      new Plugin<CommentsPluginState>({
-        key: PLUGIN_KEY,
-        state: {
-          init() {
-            return {
-              decorations: DecorationSet.empty,
-            };
-          },
-          apply(tr, state) {
-            const action = tr.getMeta(PLUGIN_KEY);
-
-            if (!tr.docChanged && !action) {
-              return state;
-            }
-
-            // only update threadPositions if the doc changed
-            const threadPositions = tr.docChanged
-              ? getUpdatedThreadPositions(tr.doc, self.markType)
-              : self.threadPositions;
-
-            if (threadPositions.size > 0 || self.threadPositions.size > 0) {
-              // small optimization; don't emit event if threadPositions before / after were both empty
-              self.threadPositions = threadPositions;
-              self.emitStateUpdate();
-            }
-
-            // update decorations if doc or selected thread changed
-            const decorations = [];
-
-            if (self.selectedThreadId) {
-              const selectedThreadPosition = threadPositions.get(
-                self.selectedThreadId,
-              );
-
-              if (selectedThreadPosition) {
-                decorations.push(
-                  Decoration.inline(
-                    selectedThreadPosition.from,
-                    selectedThreadPosition.to,
-                    {
-                      class: "bn-thread-mark-selected",
-                    },
-                  ),
-                );
-              }
-            }
-
-            return {
-              decorations: DecorationSet.create(tr.doc, decorations),
-            };
-          },
-        },
-        props: {
-          decorations(state) {
-            return (
-              PLUGIN_KEY.getState(state)?.decorations ?? DecorationSet.empty
-            );
-          },
-          /**
-           * Handle click on a thread mark and mark it as selected
-           */
-          handleClick: (view, pos, event) => {
-            if (event.button !== 0) {
-              return;
-            }
-
-            const node = view.state.doc.nodeAt(pos);
-
-            if (!node) {
-              self.selectThread(undefined);
-              return;
-            }
-
-            const markInSchema = view.state.schema.marks[markType];
-            const resolvedPos = view.state.doc.resolve(pos);
-            const commentMark = node.marks
-              .filter(
-                (mark) =>
-                  mark.type.name === markType && mark.attrs.orphan !== true,
-              )
-              .map((mark) => {
-                // get the range of this mark within the document, to check how close it is to the click position
-                const range = getMarkRange(
-                  resolvedPos,
-                  markInSchema,
-                  mark.attrs,
-                )!;
-
-                return {
-                  mark,
-                  // calculate how far the mark is from the click position
-                  distance:
-                    (Math.abs(range.from - pos) + Math.abs(range.to - pos)) / 2,
-                  // calculate the length of text the mark spans
-                  length: range.to - range.from,
-                };
-              })
-              // This allows us to not have comments which are unreachable because they are completely overlapped by other comments (issue #2073)
-              .sort((a, b) => {
-                // Find the mark which is closest to the click position
-                if (a.distance !== b.distance) {
-                  return a.distance - b.distance;
-                }
-                // Otherwise, select the mark which spans the smallest amount of text (most likely to be unreachable)
-                return a.length - b.length;
-              })[0]?.mark;
-
-            const threadId = commentMark?.attrs.threadId as string | undefined;
-            self.selectThread(threadId);
-          },
-        },
-      }),
-    );
-  }
-
-  /**
-   * Subscribe to state updates
-   */
-  public onUpdate(
-    callback: (state: {
-      pendingComment: boolean;
-      selectedThreadId: string | undefined;
-      threadPositions: Map<string, { from: number; to: number }>;
-    }) => void,
-  ) {
-    return this.on("update", callback);
-  }
-
-  /**
-   * Set the selected thread
-   */
-  public selectThread(threadId: string | undefined) {
-    if (this.selectedThreadId === threadId) {
-      return;
-    }
-    this.selectedThreadId = threadId;
-    this.emitStateUpdate();
-    this.editor.transact((tr) =>
-      tr.setMeta(PLUGIN_KEY, {
-        name: SET_SELECTED_THREAD_ID,
-      }),
-    );
-
-    if (threadId) {
-      const selectedThreadPosition = this.threadPositions.get(threadId);
-
-      if (!selectedThreadPosition) {
-        return;
-      }
-
-      // When a new thread is selected, scrolls the page to its reference text in
-      // the editor.
-      (
-        this.editor.prosemirrorView?.domAtPos(selectedThreadPosition.from)
-          .node as Element | undefined
-      )?.scrollIntoView({
-        behavior: "smooth",
-        block: "center",
-      });
-    }
-  }
-
-  /**
-   * Start a pending comment (e.g.: when clicking the "Add comment" button)
-   */
-  public startPendingComment() {
-    this.pendingComment = true;
-    this.emitStateUpdate();
-  }
-
-  /**
-   * Stop a pending comment (e.g.: user closes the comment composer)
-   */
-  public stopPendingComment() {
-    this.pendingComment = false;
-    this.emitStateUpdate();
-  }
-
-  /**
-   * Create a thread at the current selection
-   */
-  public async createThread(options: {
-    initialComment: {
-      body: CommentBody;
-      metadata?: any;
-    };
-    metadata?: any;
-  }) {
-    const thread = await this.threadStore.createThread(options);
-
-    if (this.threadStore.addThreadToDocument) {
-      // creating the mark is handled by the store
-      // this is useful if we don't have write-access to the document.
-      // We can then offload the responsibility of creating the mark to the server.
-      // (e.g.: RESTYjsThreadStore)
-      const view = this.editor.prosemirrorView!;
-      const pmSelection = view.state.selection;
-
-      const ystate = ySyncPluginKey.getState(view.state);
-
-      const selection = {
-        prosemirror: {
-          head: pmSelection.head,
-          anchor: pmSelection.anchor,
-        },
-        yjs: ystate
-          ? getRelativeSelection(ystate.binding, view.state)
-          : undefined, // if we're not using yjs
-      };
-
-      await this.threadStore.addThreadToDocument({
-        threadId: thread.id,
-        selection,
-      });
-    } else {
-      // we create the mark directly in the document
-      this.editor._tiptapEditor.commands.setMark(this.markType, {
-        orphan: false,
-        threadId: thread.id,
-      });
-    }
-  }
-}
diff --git a/packages/core/src/extensions/DropCursor/DropCursor.ts b/packages/core/src/extensions/DropCursor/DropCursor.ts
new file mode 100644
index 0000000000..97e0bed1bb
--- /dev/null
+++ b/packages/core/src/extensions/DropCursor/DropCursor.ts
@@ -0,0 +1,15 @@
+import { dropCursor } from "prosemirror-dropcursor";
+import { createExtension } from "../../editor/BlockNoteExtension.js";
+
+export const DropCursor = createExtension((editor, options) => {
+  return {
+    key: "dropCursor",
+    prosemirrorPlugins: [
+      (options.dropCursor ?? dropCursor)({
+        width: 5,
+        color: "#ddeeff",
+        editor: editor,
+      }),
+    ],
+  } as const;
+});
diff --git a/packages/core/src/extensions/FilePanel/FilePanel.ts b/packages/core/src/extensions/FilePanel/FilePanel.ts
new file mode 100644
index 0000000000..7897dddb48
--- /dev/null
+++ b/packages/core/src/extensions/FilePanel/FilePanel.ts
@@ -0,0 +1,43 @@
+import {
+  createExtension,
+  createStore,
+} from "../../editor/BlockNoteExtension.js";
+
+export const FilePanel = createExtension((editor) => {
+  const store = createStore({
+    blockId: undefined as string | undefined,
+  });
+
+  function closeMenu() {
+    store.setState({
+      blockId: undefined,
+    });
+  }
+
+  return {
+    key: "filePanel",
+    store,
+    mount({ signal }) {
+      // reset the menu when the document changes
+      const unsubscribeOnChange = editor.onChange(
+        closeMenu,
+        // don't trigger the callback if the changes are caused by a remote user
+        false,
+      );
+
+      // reset the menu when the selection changes
+      const unsubscribeOnSelectionChange = editor.onSelectionChange(closeMenu);
+
+      signal.addEventListener("abort", () => {
+        unsubscribeOnChange();
+        unsubscribeOnSelectionChange();
+      });
+    },
+    closeMenu,
+    showMenu(blockId: string) {
+      store.setState({
+        blockId,
+      });
+    },
+  } as const;
+});
diff --git a/packages/core/src/extensions/FilePanel/FilePanelPlugin.ts b/packages/core/src/extensions/FilePanel/FilePanelPlugin.ts
deleted file mode 100644
index c5819a1c5b..0000000000
--- a/packages/core/src/extensions/FilePanel/FilePanelPlugin.ts
+++ /dev/null
@@ -1,206 +0,0 @@
-import { EditorState, Plugin, PluginKey, PluginView } from "prosemirror-state";
-import { EditorView } from "prosemirror-view";
-
-import { ySyncPluginKey } from "y-prosemirror";
-import type { BlockNoteEditor } from "../../editor/BlockNoteEditor.js";
-import { BlockNoteExtension } from "../../editor/BlockNoteExtension.js";
-import { UiElementPosition } from "../../extensions-shared/UiElementPosition.js";
-import type {
-  BlockFromConfig,
-  InlineContentSchema,
-  StyleSchema,
-} from "../../schema/index.js";
-
-export type FilePanelState<
-  I extends InlineContentSchema,
-  S extends StyleSchema,
-> = UiElementPosition & {
-  // TODO: This typing is not quite right (children should be from BSchema)
-  block: BlockFromConfig<any, I, S>;
-};
-
-export class FilePanelView<I extends InlineContentSchema, S extends StyleSchema>
-  implements PluginView
-{
-  public state?: FilePanelState<I, S>;
-  public emitUpdate: () => void;
-
-  constructor(
-    private readonly editor: BlockNoteEditor<Record<string, any>, I, S>,
-    private readonly pluginKey: PluginKey<FilePanelState<I, S>>,
-    private readonly pmView: EditorView,
-    emitUpdate: (state: FilePanelState<I, S>) => void,
-  ) {
-    this.emitUpdate = () => {
-      if (!this.state) {
-        throw new Error("Attempting to update uninitialized file panel");
-      }
-
-      emitUpdate(this.state);
-    };
-
-    pmView.dom.addEventListener("mousedown", this.mouseDownHandler);
-    pmView.dom.addEventListener("dragstart", this.dragstartHandler);
-
-    // Setting capture=true ensures that any parent container of the editor that
-    // gets scrolled will trigger the scroll event. Scroll events do not bubble
-    // and so won't propagate to the document by default.
-    pmView.root.addEventListener("scroll", this.scrollHandler, true);
-  }
-
-  mouseDownHandler = () => {
-    if (this.state?.show) {
-      this.state.show = false;
-      this.emitUpdate();
-    }
-  };
-
-  // For dragging the whole editor.
-  dragstartHandler = () => {
-    if (this.state?.show) {
-      this.state.show = false;
-      this.emitUpdate();
-    }
-  };
-
-  scrollHandler = () => {
-    if (this.state?.show) {
-      const blockElement = this.pmView.root.querySelector(
-        `[data-node-type="blockContainer"][data-id="${this.state.block.id}"]`,
-      );
-      if (!blockElement) {
-        return;
-      }
-      this.state.referencePos = blockElement.getBoundingClientRect();
-      this.emitUpdate();
-    }
-  };
-
-  update(view: EditorView, prevState: EditorState) {
-    const pluginState = this.pluginKey.getState(view.state);
-    const prevPluginState = this.pluginKey.getState(prevState);
-
-    if (!this.state?.show && pluginState?.block && this.editor.isEditable) {
-      const blockElement = this.pmView.root.querySelector(
-        `[data-node-type="blockContainer"][data-id="${pluginState.block.id}"]`,
-      );
-      if (!blockElement) {
-        return;
-      }
-      this.state = {
-        show: true,
-        referencePos: blockElement.getBoundingClientRect(),
-        block: pluginState.block,
-      };
-
-      this.emitUpdate();
-
-      return;
-    }
-
-    const isOpening = pluginState?.block && !prevPluginState?.block;
-    const isClosing = !pluginState?.block && prevPluginState?.block;
-    if (isOpening && this.state && !this.state.show) {
-      this.state.show = true;
-      this.emitUpdate();
-    }
-    if (isClosing && this.state?.show) {
-      this.state.show = false;
-      this.emitUpdate();
-    }
-  }
-
-  closeMenu = () => {
-    if (this.state?.show) {
-      this.state.show = false;
-      this.emitUpdate();
-    }
-  };
-
-  destroy() {
-    this.pmView.dom.removeEventListener("mousedown", this.mouseDownHandler);
-
-    this.pmView.dom.removeEventListener("dragstart", this.dragstartHandler);
-
-    this.pmView.root.removeEventListener("scroll", this.scrollHandler, true);
-  }
-}
-
-const filePanelPluginKey = new PluginKey<FilePanelState<any, any>>(
-  "FilePanelPlugin",
-);
-
-export class FilePanelProsemirrorPlugin<
-  I extends InlineContentSchema,
-  S extends StyleSchema,
-> extends BlockNoteExtension {
-  public static key() {
-    return "filePanel";
-  }
-
-  private view: FilePanelView<I, S> | undefined;
-
-  constructor(editor: BlockNoteEditor<Record<string, any>, I, S>) {
-    super();
-    this.addProsemirrorPlugin(
-      new Plugin<{
-        block: BlockFromConfig<any, I, S> | undefined;
-      }>({
-        key: filePanelPluginKey,
-        view: (editorView) => {
-          this.view = new FilePanelView<I, S>(
-            editor,
-            filePanelPluginKey as any,
-            editorView,
-            (state) => {
-              this.emit("update", state);
-            },
-          );
-          return this.view;
-        },
-        props: {
-          handleKeyDown: (_view, event: KeyboardEvent) => {
-            if (event.key === "Escape" && this.shown) {
-              this.view?.closeMenu();
-              return true;
-            }
-            return false;
-          },
-        },
-        state: {
-          init: () => {
-            return {
-              block: undefined,
-            };
-          },
-          apply: (transaction, prev) => {
-            const state: FilePanelState<I, S> | undefined =
-              transaction.getMeta(filePanelPluginKey);
-
-            if (state) {
-              return state;
-            }
-
-            if (
-              !transaction.getMeta(ySyncPluginKey) &&
-              (transaction.selectionSet || transaction.docChanged)
-            ) {
-              return { block: undefined };
-            }
-            return prev;
-          },
-        },
-      }),
-    );
-  }
-
-  public get shown() {
-    return this.view?.state?.show || false;
-  }
-
-  public onUpdate(callback: (state: FilePanelState<I, S>) => void) {
-    return this.on("update", callback);
-  }
-
-  public closeMenu = () => this.view?.closeMenu();
-}
diff --git a/packages/core/src/extensions/FormattingToolbar/FormattingToolbar.ts b/packages/core/src/extensions/FormattingToolbar/FormattingToolbar.ts
new file mode 100644
index 0000000000..24f82a0346
--- /dev/null
+++ b/packages/core/src/extensions/FormattingToolbar/FormattingToolbar.ts
@@ -0,0 +1,100 @@
+import {
+  createExtension,
+  createStore,
+} from "../../editor/BlockNoteExtension.js";
+import { TextSelection } from "prosemirror-state";
+
+export const FormattingToolbar = createExtension((editor) => {
+  const store = createStore(false);
+
+  const shouldShow = () => {
+    return editor.transact((tr) => {
+      // if the selection is empty, or not a text selection, we should not show the toolbar
+      if (tr.selection.empty || !(tr.selection instanceof TextSelection)) {
+        return false;
+      }
+
+      // We don't want to show the toolbar if the selection has code in it
+      let spansCode = false;
+      // search the content of the selection to see if it ever spans a node with a code spec
+      tr.selection.content().content.descendants((node) => {
+        if (node.type.spec.code) {
+          spansCode = true;
+        }
+        return !spansCode; // keep descending if we haven't found a code block
+      });
+
+      // if the selection is in a code block, we should not show the toolbar
+      if (spansCode) {
+        return false;
+      }
+
+      // otherwise, we should show the toolbar at the current selection
+      return true;
+    });
+  };
+
+  return {
+    key: "formattingToolbar",
+    store: store,
+    mount({ dom, signal }) {
+      /**
+       * We want to mimic the Notion behavior of not showing the toolbar while the user is holding down the mouse button (to create a selection)
+       */
+      let preventShowWhileMouseDown = false;
+
+      const unsubscribeOnChange = editor.onChange(() => {
+        if (preventShowWhileMouseDown) {
+          return;
+        }
+        // re-evaluate whether the toolbar should be shown
+        store.setState(shouldShow());
+      });
+      const unsubscribeOnSelectionChange = editor.onSelectionChange(() => {
+        if (preventShowWhileMouseDown) {
+          return;
+        }
+        // re-evaluate whether the toolbar should be shown
+        store.setState(shouldShow());
+      });
+
+      // To mimic Notion's behavior, we listen to the mouse down event to set the `preventShowWhileMouseDown` flag
+      dom.addEventListener(
+        "pointerdown",
+        () => {
+          preventShowWhileMouseDown = true;
+          store.setState(false);
+        },
+        { signal },
+      );
+      // To mimic Notion's behavior, we listen to the mouse up event to reset the `preventShowWhileMouseDown` flag and show the toolbar (if it should)
+      editor.prosemirrorView.root.addEventListener(
+        "pointerup",
+        () => {
+          preventShowWhileMouseDown = false;
+          // We only want to re-show the toolbar if the mouse made the selection
+          if (editor.isFocused()) {
+            store.setState(shouldShow());
+          }
+        },
+        { signal, capture: true },
+      );
+      // If the pointer gets cancelled, we don't want to be stuck in the `preventShowWhileMouseDown` state
+      dom.addEventListener(
+        "pointercancel",
+        () => {
+          preventShowWhileMouseDown = false;
+        },
+        {
+          signal,
+          capture: true,
+        },
+      );
+
+      signal.addEventListener("abort", () => {
+        unsubscribeOnChange();
+        unsubscribeOnSelectionChange();
+      });
+    },
+  } as const;
+});
diff --git a/packages/core/src/extensions/FormattingToolbar/FormattingToolbarPlugin.ts b/packages/core/src/extensions/FormattingToolbar/FormattingToolbarPlugin.ts
deleted file mode 100644
index 5bfa0f5dc8..0000000000
--- a/packages/core/src/extensions/FormattingToolbar/FormattingToolbarPlugin.ts
+++ /dev/null
@@ -1,363 +0,0 @@
-import { isNodeSelection, isTextSelection, posToDOMRect } from "@tiptap/core";
-import {
-  EditorState,
-  Plugin,
-  PluginKey,
-  PluginView,
-  TextSelection,
-} from "prosemirror-state";
-import { EditorView } from "prosemirror-view";
-
-import type { BlockNoteEditor } from "../../editor/BlockNoteEditor.js";
-import { BlockNoteExtension } from "../../editor/BlockNoteExtension.js";
-import { UiElementPosition } from "../../extensions-shared/UiElementPosition.js";
-import {
-  BlockSchema,
-  InlineContentSchema,
-  StyleSchema,
-} from "../../schema/index.js";
-
-export type FormattingToolbarState = UiElementPosition;
-
-export class FormattingToolbarView implements PluginView {
-  public state?: FormattingToolbarState;
-  public emitUpdate: () => void;
-
-  public preventHide = false;
-  public preventShow = false;
-
-  public shouldShow: (props: {
-    view: EditorView;
-    state: EditorState;
-    from: number;
-    to: number;
-  }) => boolean = ({ view, state, from, to }) => {
-    const { doc, selection } = state;
-    const { empty } = selection;
-
-    // Sometime check for `empty` is not enough.
-    // Doubleclick an empty paragraph returns a node size of 2.
-    // So we check also for an empty text size.
-    const isEmptyTextBlock =
-      !doc.textBetween(from, to).length && isTextSelection(state.selection);
-
-    // Don't show toolbar inside code blocks
-    if (
-      selection.$from.parent.type.spec.code ||
-      (isNodeSelection(selection) && selection.node.type.spec.code)
-    ) {
-      return false;
-    }
-
-    if (empty || isEmptyTextBlock) {
-      return false;
-    }
-
-    const focusedElement = document.activeElement;
-    if (!this.isElementWithinEditorWrapper(focusedElement) && view.editable) {
-      // editable editors must have focus for the toolbar to show
-      return false;
-    }
-    return true;
-  };
-
-  constructor(
-    private readonly editor: BlockNoteEditor<
-      BlockSchema,
-      InlineContentSchema,
-      StyleSchema
-    >,
-    private readonly pmView: EditorView,
-    emitUpdate: (state: FormattingToolbarState) => void,
-  ) {
-    this.emitUpdate = () => {
-      if (!this.state) {
-        throw new Error(
-          "Attempting to update uninitialized formatting toolbar",
-        );
-      }
-
-      emitUpdate(this.state);
-    };
-
-    pmView.dom.addEventListener("mousedown", this.viewMousedownHandler);
-    pmView.root.addEventListener("mouseup", this.mouseupHandler);
-    pmView.dom.addEventListener("dragstart", this.dragHandler);
-    pmView.dom.addEventListener("dragover", this.dragHandler);
-    pmView.dom.addEventListener("blur", this.blurHandler);
-
-    // Setting capture=true ensures that any parent container of the editor that
-    // gets scrolled will trigger the scroll event. Scroll events do not bubble
-    // and so won't propagate to the document by default.
-    pmView.root.addEventListener("scroll", this.scrollHandler, true);
-  }
-
-  blurHandler = (event: FocusEvent) => {
-    if (this.preventHide) {
-      this.preventHide = false;
-
-      return;
-    }
-
-    const editorWrapper = this.pmView.dom.parentElement!;
-
-    // Checks if the focus is moving to an element outside the editor. If it is,
-    // the toolbar is hidden.
-    if (
-      // An element is clicked.
-      event &&
-      event.relatedTarget &&
-      // Element is inside the editor.
-      (editorWrapper === (event.relatedTarget as Node) ||
-        editorWrapper.contains(event.relatedTarget as Node) ||
-        (event.relatedTarget as HTMLElement).matches(
-          ".bn-ui-container, .bn-ui-container *",
-        ))
-    ) {
-      return;
-    }
-
-    if (this.state?.show) {
-      this.state.show = false;
-      this.emitUpdate();
-    }
-  };
-
-  isElementWithinEditorWrapper = (element: Node | null) => {
-    if (!element) {
-      return false;
-    }
-    const editorWrapper = this.pmView.dom.parentElement!;
-    if (!editorWrapper) {
-      return false;
-    }
-
-    return editorWrapper.contains(element);
-  };
-
-  viewMousedownHandler = (e: MouseEvent) => {
-    if (
-      !this.isElementWithinEditorWrapper(e.target as Node) ||
-      e.button === 0
-    ) {
-      this.preventShow = true;
-    }
-  };
-
-  mouseupHandler = () => {
-    if (this.preventShow) {
-      this.preventShow = false;
-      setTimeout(() => this.update(this.pmView));
-    }
-  };
-
-  // For dragging the whole editor.
-  dragHandler = () => {
-    if (this.state?.show) {
-      this.state.show = false;
-      this.emitUpdate();
-    }
-  };
-
-  scrollHandler = () => {
-    if (this.state?.show) {
-      this.state.referencePos = this.getSelectionBoundingBox();
-      this.emitUpdate();
-    }
-  };
-
-  update(view: EditorView, oldState?: EditorState) {
-    // Delays the update to handle edge case with drag and drop, where the view
-    // is blurred asynchronously and happens only after the state update.
-    // Wrapping in a setTimeout gives enough time to wait for the blur event to
-    // occur before updating the toolbar.
-    const { state, composing } = view;
-    const { doc, selection } = state;
-    const isSame =
-      oldState && oldState.doc.eq(doc) && oldState.selection.eq(selection);
-
-    if (composing || isSame) {
-      return;
-    }
-
-    // support for CellSelections
-    const { ranges } = selection;
-    const from = Math.min(...ranges.map((range) => range.$from.pos));
-    const to = Math.max(...ranges.map((range) => range.$to.pos));
-
-    const shouldShow = this.shouldShow({
-      view,
-      state,
-      from,
-      to,
-    });
-
-    // in jsdom, Range.prototype.getClientRects is not implemented,
-    // this would cause `getSelectionBoundingBox` to fail
-    // we can just ignore jsdom for now and not show the toolbar
-    const jsdom = typeof Range.prototype.getClientRects === "undefined";
-
-    // Checks if menu should be shown/updated.
-    if (!this.preventShow && (shouldShow || this.preventHide) && !jsdom) {
-      // Unlike other UI elements, we don't prevent the formatting toolbar from
-      // showing when the editor is not editable. This is because some buttons,
-      // e.g. the download file button, should still be accessible. Therefore,
-      // logic for hiding when the editor is non-editable is handled
-      // individually in each button.
-      const newReferencePos = this.getSelectionBoundingBox();
-
-      // Workaround to ensure the correct reference position when rendering
-      // React components. Without this, e.g. updating styles on React inline
-      // content causes the formatting toolbar to be in the wrong place. We
-      // know the component has not yet rendered if the reference position has
-      // zero dimensions.
-      if (newReferencePos.height === 0 && newReferencePos.width === 0) {
-        // Updates the reference position again following the render.
-        queueMicrotask(() => {
-          const nextState = {
-            show: true,
-            referencePos: this.getSelectionBoundingBox(),
-          };
-
-          this.state = nextState;
-          this.emitUpdate();
-
-          // For some reason, while the selection doesn't actually change and
-          // remains correct, it visually appears to be collapsed. This forces
-          // a ProseMirror view update, which fixes the issue.
-          view.dispatch(
-            view.state.tr.setSelection(
-              TextSelection.create(
-                view.state.doc,
-                view.state.selection.from + 1,
-                view.state.selection.to,
-              ),
-            ),
-          );
-          // 2 separate `dispatch` calls are needed, else ProseMirror realizes
-          // that the transaction is a no-op and doesn't update the view.
-          view.dispatch(
-            view.state.tr.setSelection(
-              TextSelection.create(
-                view.state.doc,
-                view.state.selection.from - 1,
-                view.state.selection.to,
-              ),
-            ),
-          );
-        });
-
-        return;
-      }
-
-      const nextState = {
-        show: true,
-        referencePos: this.getSelectionBoundingBox(),
-      };
-
-      if (
-        nextState.show !== this.state?.show ||
-        nextState.referencePos.toJSON() !== this.state?.referencePos.toJSON()
-      ) {
-        this.state = nextState;
-        this.emitUpdate();
-      }
-
-      return;
-    }
-
-    // Checks if menu should be hidden.
-    if (
-      this.state?.show &&
-      !this.preventHide &&
-      (!shouldShow || this.preventShow || !this.editor.isEditable)
-    ) {
-      this.state.show = false;
-      this.emitUpdate();
-
-      return;
-    }
-  }
-
-  destroy() {
-    this.pmView.dom.removeEventListener("mousedown", this.viewMousedownHandler);
-    this.pmView.root.removeEventListener("mouseup", this.mouseupHandler);
-    this.pmView.dom.removeEventListener("dragstart", this.dragHandler);
-    this.pmView.dom.removeEventListener("dragover", this.dragHandler);
-    this.pmView.dom.removeEventListener("blur", this.blurHandler);
-
-    this.pmView.root.removeEventListener("scroll", this.scrollHandler, true);
-  }
-
-  closeMenu = () => {
-    if (this.state?.show) {
-      this.state.show = false;
-      this.emitUpdate();
-    }
-  };
-
-  getSelectionBoundingBox() {
-    const { state } = this.pmView;
-    const { selection } = state;
-
-    // support for CellSelections
-    const { ranges } = selection;
-    const from = Math.min(...ranges.map((range) => range.$from.pos));
-    const to = Math.max(...ranges.map((range) => range.$to.pos));
-
-    if (isNodeSelection(selection)) {
-      const node = this.pmView.nodeDOM(from) as HTMLElement;
-      if (node) {
-        return node.getBoundingClientRect();
-      }
-    }
-
-    return posToDOMRect(this.pmView, from, to);
-  }
-}
-
-export const formattingToolbarPluginKey = new PluginKey(
-  "FormattingToolbarPlugin",
-);
-
-export class FormattingToolbarProsemirrorPlugin extends BlockNoteExtension {
-  public static key() {
-    return "formattingToolbar";
-  }
-
-  private view: FormattingToolbarView | undefined;
-
-  constructor(editor: BlockNoteEditor<any, any, any>) {
-    super();
-    this.addProsemirrorPlugin(
-      new Plugin({
-        key: formattingToolbarPluginKey,
-        view: (editorView) => {
-          this.view = new FormattingToolbarView(editor, editorView, (state) => {
-            this.emit("update", state);
-          });
-          return this.view;
-        },
-        props: {
-          handleKeyDown: (_view, event: KeyboardEvent) => {
-            if (event.key === "Escape" && this.shown) {
-              this.view!.closeMenu();
-              return true;
-            }
-            return false;
-          },
-        },
-      }),
-    );
-  }
-
-  public get shown() {
-    return this.view?.state?.show || false;
-  }
-
-  public onUpdate(callback: (state: FormattingToolbarState) => void) {
-    return this.on("update", callback);
-  }
-
-  public closeMenu = () => this.view!.closeMenu();
-}
diff --git a/packages/core/src/extensions/History/History.ts b/packages/core/src/extensions/History/History.ts
new file mode 100644
index 0000000000..08e82daa58
--- /dev/null
+++ b/packages/core/src/extensions/History/History.ts
@@ -0,0 +1,15 @@
+import { history, redo, undo } from "@tiptap/pm/history";
+import { createExtension } from "../../editor/BlockNoteExtension.js";
+
+export const History = createExtension((_editor, options) => {
+  if (options.collaboration) {
+    return;
+  }
+
+  return {
+    key: "history",
+    prosemirrorPlugins: [history()],
+    undoCommand: undo,
+    redoCommand: redo,
+  } as const;
+});
diff --git a/packages/core/src/extensions/LinkToolbar/LinkToolbar.ts b/packages/core/src/extensions/LinkToolbar/LinkToolbar.ts
new file mode 100644
index 0000000000..d55fb6122c
--- /dev/null
+++ b/packages/core/src/extensions/LinkToolbar/LinkToolbar.ts
@@ -0,0 +1,114 @@
+import { getMarkRange, posToDOMRect } from "@tiptap/core";
+import { createExtension } from "../../editor/BlockNoteExtension.js";
+import { getPmSchema } from "../../api/pmUtil.js";
+
+export const LinkToolbar = createExtension((editor) => {
+  function getLinkElementAtPos(pos: number) {
+    let currentNode = editor.prosemirrorView.nodeDOM(pos);
+    while (currentNode && currentNode.parentElement) {
+      if (currentNode.nodeName === "A") {
+        return currentNode as HTMLAnchorElement;
+      }
+      currentNode = currentNode.parentElement;
+    }
+    return null;
+  }
+
+  function getMarkAtPos(pos: number, markType: string) {
+    return editor.transact((tr) => {
+      const resolvedPos = tr.doc.resolve(pos);
+      const mark = resolvedPos
+        .marks()
+        .find((mark) => mark.type.name === markType);
+
+      if (!mark) {
+        return;
+      }
+
+      const markRange = getMarkRange(resolvedPos, mark.type);
+      if (!markRange) {
+        return;
+      }
+
+      return {
+        range: markRange,
+        mark,
+        get text() {
+          return tr.doc.textBetween(markRange.from, markRange.to);
+        },
+        get position() {
+          // to minimize re-renders, we convert to JSON, which is the same shape anyway
+          return posToDOMRect(
+            editor.prosemirrorView,
+            markRange.from,
+            markRange.to,
+          ).toJSON() as DOMRect;
+        },
+      };
+    });
+  }
+
+  function getLinkAtSelection() {
+    return editor.transact((tr) => {
+      const selection = tr.selection;
+      return getMarkAtPos(selection.anchor, "link");
+    });
+  }
+
+  return {
+    key: "linkToolbar",
+
+    getLinkAtSelection,
+    getLinkElementAtPos,
+    getMarkAtPos,
+
+    getLinkAtElement(element: HTMLElement) {
+      return editor.transact(() => {
+        const posAtElement = editor.prosemirrorView.posAtDOM(element, 0) + 1;
+        return getMarkAtPos(posAtElement, "link");
+      });
+    },
+
+    editLink(
+      url: string,
+      text: string,
+      position = editor.transact((tr) => tr.selection.anchor),
+    ) {
+      editor.transact((tr) => {
+        const pmSchema = getPmSchema(tr);
+        const { range } = getMarkAtPos(position + 1, "link") || {
+          range: {
+            from: tr.selection.from,
+            to: tr.selection.to,
+          },
+        };
+        if (!range) {
+          return;
+        }
+        tr.insertText(text, range.from, range.to);
+        tr.addMark(
+          range.from,
+          range.from + text.length,
+          pmSchema.mark("link", { href: url }),
+        );
+      });
+      editor.prosemirrorView.focus();
+    },
+
+    deleteLink() {
+      editor.transact((tr) => {
+        const pmSchema = getPmSchema(tr);
+        const range = getLinkAtSelection()?.range;
+        if (!range) {
+          return;
+        }
+
+        tr.removeMark(range.from, range.to, pmSchema.mark("link")).setMeta(
+          "preventAutolink",
+          true,
+        );
+      });
+      editor.prosemirrorView.focus();
+    },
+  } as const;
+});
diff --git a/packages/core/src/extensions/LinkToolbar/LinkToolbarPlugin.ts b/packages/core/src/extensions/LinkToolbar/LinkToolbarPlugin.ts
deleted file mode 100644
index a6b38e7f6f..0000000000
--- a/packages/core/src/extensions/LinkToolbar/LinkToolbarPlugin.ts
+++ /dev/null
@@ -1,380 +0,0 @@
-import { getMarkRange, posToDOMRect, Range } from "@tiptap/core";
-
-import { EditorView } from "@tiptap/pm/view";
-import { Mark } from "prosemirror-model";
-import { EditorState, Plugin, PluginKey, PluginView } from "prosemirror-state";
-
-import { getPmSchema } from "../../api/pmUtil.js";
-import type { BlockNoteEditor } from "../../editor/BlockNoteEditor.js";
-import { BlockNoteExtension } from "../../editor/BlockNoteExtension.js";
-import { UiElementPosition } from "../../extensions-shared/UiElementPosition.js";
-import {
-  BlockSchema,
-  InlineContentSchema,
-  StyleSchema,
-} from "../../schema/index.js";
-
-export type LinkToolbarState = UiElementPosition & {
-  // The hovered link's URL, and the text it's displayed with in the
-  // editor.
-  url: string;
-  text: string;
-};
-
-class LinkToolbarView implements PluginView {
-  public state?: LinkToolbarState;
-  public emitUpdate: () => void;
-
-  menuUpdateTimer: ReturnType<typeof setTimeout> | undefined;
-  startMenuUpdateTimer: () => void;
-  stopMenuUpdateTimer: () => void;
-
-  mouseHoveredLinkMark: Mark | undefined;
-  mouseHoveredLinkMarkRange: Range | undefined;
-
-  keyboardHoveredLinkMark: Mark | undefined;
-  keyboardHoveredLinkMarkRange: Range | undefined;
-
-  linkMark: Mark | undefined;
-  linkMarkRange: Range | undefined;
-
-  constructor(
-    private readonly editor: BlockNoteEditor<any, any, any>,
-    private readonly pmView: EditorView,
-    emitUpdate: (state: LinkToolbarState) => void,
-  ) {
-    this.emitUpdate = () => {
-      if (!this.state) {
-        throw new Error("Attempting to update uninitialized link toolbar");
-      }
-
-      emitUpdate(this.state);
-    };
-
-    this.startMenuUpdateTimer = () => {
-      this.menuUpdateTimer = setTimeout(() => {
-        this.update(this.pmView, undefined, true);
-      }, 250);
-    };
-
-    this.stopMenuUpdateTimer = () => {
-      if (this.menuUpdateTimer) {
-        clearTimeout(this.menuUpdateTimer);
-        this.menuUpdateTimer = undefined;
-      }
-
-      return false;
-    };
-
-    this.pmView.dom.addEventListener("mouseover", this.mouseOverHandler);
-    this.pmView.root.addEventListener(
-      "click",
-      this.clickHandler as EventListener,
-      true,
-    );
-
-    // Setting capture=true ensures that any parent container of the editor that
-    // gets scrolled will trigger the scroll event. Scroll events do not bubble
-    // and so won't propagate to the document by default.
-    this.pmView.root.addEventListener("scroll", this.scrollHandler, true);
-  }
-
-  mouseOverHandler = (event: MouseEvent) => {
-    // Resets the link mark currently hovered by the mouse cursor.
-    this.mouseHoveredLinkMark = undefined;
-    this.mouseHoveredLinkMarkRange = undefined;
-
-    this.stopMenuUpdateTimer();
-
-    if (
-      event.target instanceof HTMLAnchorElement &&
-      event.target.nodeName === "A"
-    ) {
-      // Finds link mark at the hovered element's position to update mouseHoveredLinkMark and
-      // mouseHoveredLinkMarkRange.
-      const hoveredLinkElement = event.target;
-      const posInHoveredLinkMark =
-        this.pmView.posAtDOM(hoveredLinkElement, 0) + 1;
-      const resolvedPosInHoveredLinkMark =
-        this.pmView.state.doc.resolve(posInHoveredLinkMark);
-      const marksAtPos = resolvedPosInHoveredLinkMark.marks();
-
-      for (const mark of marksAtPos) {
-        if (
-          mark.type.name === this.pmView.state.schema.mark("link").type.name
-        ) {
-          this.mouseHoveredLinkMark = mark;
-          this.mouseHoveredLinkMarkRange =
-            getMarkRange(resolvedPosInHoveredLinkMark, mark.type, mark.attrs) ||
-            undefined;
-
-          break;
-        }
-      }
-    }
-
-    this.startMenuUpdateTimer();
-
-    return false;
-  };
-
-  clickHandler = (event: MouseEvent) => {
-    const editorWrapper = this.pmView.dom.parentElement!;
-
-    if (
-      // Toolbar is open.
-      this.linkMark &&
-      // An element is clicked.
-      event &&
-      event.target &&
-      // The clicked element is not the editor.
-      !(
-        editorWrapper === (event.target as Node) ||
-        editorWrapper.contains(event.target as Node)
-      )
-    ) {
-      if (this.state?.show) {
-        this.state.show = false;
-        this.emitUpdate();
-      }
-    }
-  };
-
-  scrollHandler = () => {
-    if (this.linkMark !== undefined) {
-      if (this.state?.show) {
-        this.state.referencePos = posToDOMRect(
-          this.pmView,
-          this.linkMarkRange!.from,
-          this.linkMarkRange!.to,
-        );
-        this.emitUpdate();
-      }
-    }
-  };
-
-  editLink(url: string, text: string) {
-    this.editor.transact((tr) => {
-      const pmSchema = getPmSchema(tr);
-      tr.insertText(text, this.linkMarkRange!.from, this.linkMarkRange!.to);
-      tr.addMark(
-        this.linkMarkRange!.from,
-        this.linkMarkRange!.from + text.length,
-        pmSchema.mark("link", { href: url }),
-      );
-    });
-    this.pmView.focus();
-
-    if (this.state?.show) {
-      this.state.show = false;
-      this.emitUpdate();
-    }
-  }
-
-  deleteLink() {
-    this.editor.transact((tr) =>
-      tr
-        .removeMark(
-          this.linkMarkRange!.from,
-          this.linkMarkRange!.to,
-          this.linkMark!.type,
-        )
-        .setMeta("preventAutolink", true),
-    );
-    this.pmView.focus();
-
-    if (this.state?.show) {
-      this.state.show = false;
-      this.emitUpdate();
-    }
-  }
-
-  update(view: EditorView, oldState?: EditorState, fromMouseOver = false) {
-    const { state } = view;
-
-    const isSame =
-      oldState &&
-      oldState.selection.from === state.selection.from &&
-      oldState.selection.to === state.selection.to;
-
-    if (isSame || !this.pmView.hasFocus()) {
-      return;
-    }
-
-    // Saves the currently hovered link mark before it's updated.
-    const prevLinkMark = this.linkMark;
-
-    // Resets the currently hovered link mark.
-    this.linkMark = undefined;
-    this.linkMarkRange = undefined;
-
-    // Resets the link mark currently hovered by the keyboard cursor.
-    this.keyboardHoveredLinkMark = undefined;
-    this.keyboardHoveredLinkMarkRange = undefined;
-
-    // Finds link mark at the editor selection's position to update keyboardHoveredLinkMark and
-    // keyboardHoveredLinkMarkRange.
-    if (this.pmView.state.selection.empty) {
-      const marksAtPos = this.pmView.state.selection.$from.marks();
-
-      for (const mark of marksAtPos) {
-        if (
-          mark.type.name === this.pmView.state.schema.mark("link").type.name
-        ) {
-          this.keyboardHoveredLinkMark = mark;
-          this.keyboardHoveredLinkMarkRange =
-            getMarkRange(
-              this.pmView.state.selection.$from,
-              mark.type,
-              mark.attrs,
-            ) || undefined;
-
-          break;
-        }
-      }
-    }
-
-    if (this.mouseHoveredLinkMark && fromMouseOver) {
-      this.linkMark = this.mouseHoveredLinkMark;
-      this.linkMarkRange = this.mouseHoveredLinkMarkRange;
-    }
-
-    // Keyboard cursor position takes precedence over mouse hovered link.
-    if (this.keyboardHoveredLinkMark) {
-      this.linkMark = this.keyboardHoveredLinkMark;
-      this.linkMarkRange = this.keyboardHoveredLinkMarkRange;
-    }
-
-    if (this.linkMark && this.editor.isEditable) {
-      this.state = {
-        show: true,
-        referencePos: posToDOMRect(
-          this.pmView,
-          this.linkMarkRange!.from,
-          this.linkMarkRange!.to,
-        ),
-        url: this.linkMark!.attrs.href,
-        text: this.pmView.state.doc.textBetween(
-          this.linkMarkRange!.from,
-          this.linkMarkRange!.to,
-        ),
-      };
-      this.emitUpdate();
-
-      return;
-    }
-
-    // Hides menu.
-    if (
-      this.state?.show &&
-      prevLinkMark &&
-      (!this.linkMark || !this.editor.isEditable)
-    ) {
-      this.state.show = false;
-      this.emitUpdate();
-
-      return;
-    }
-  }
-
-  closeMenu = () => {
-    if (this.state?.show) {
-      this.state.show = false;
-      this.emitUpdate();
-    }
-  };
-
-  destroy() {
-    this.pmView.dom.removeEventListener("mouseover", this.mouseOverHandler);
-    this.pmView.root.removeEventListener("scroll", this.scrollHandler, true);
-    this.pmView.root.removeEventListener(
-      "click",
-      this.clickHandler as EventListener,
-      true,
-    );
-  }
-}
-
-export const linkToolbarPluginKey = new PluginKey("LinkToolbarPlugin");
-
-export class LinkToolbarProsemirrorPlugin<
-  BSchema extends BlockSchema,
-  I extends InlineContentSchema,
-  S extends StyleSchema,
-> extends BlockNoteExtension {
-  public static key() {
-    return "linkToolbar";
-  }
-
-  private view: LinkToolbarView | undefined;
-
-  constructor(editor: BlockNoteEditor<BSchema, I, S>) {
-    super();
-    this.addProsemirrorPlugin(
-      new Plugin({
-        key: linkToolbarPluginKey,
-        view: (editorView) => {
-          this.view = new LinkToolbarView(editor, editorView, (state) => {
-            this.emit("update", state);
-          });
-          return this.view;
-        },
-        props: {
-          handleKeyDown: (_view, event: KeyboardEvent) => {
-            if (event.key === "Escape" && this.shown) {
-              this.view!.closeMenu();
-              return true;
-            }
-            return false;
-          },
-        },
-      }),
-    );
-  }
-
-  public onUpdate(callback: (state: LinkToolbarState) => void) {
-    return this.on("update", callback);
-  }
-
-  /**
-   * Edit the currently hovered link.
-   */
-  public editLink = (url: string, text: string) => {
-    this.view!.editLink(url, text);
-  };
-
-  /**
-   * Delete the currently hovered link.
-   */
-  public deleteLink = () => {
-    this.view!.deleteLink();
-  };
-
-  /**
-   * When hovering on/off links using the mouse cursor, the link toolbar will
-   * open & close with a delay.
-   *
-   * This function starts the delay timer, and should be used for when the mouse
-   * cursor enters the link toolbar.
-   */
-  public startHideTimer = () => {
-    this.view!.startMenuUpdateTimer();
-  };
-
-  /**
-   * When hovering on/off links using the mouse cursor, the link toolbar will
-   * open & close with a delay.
-   *
-   * This function stops the delay timer, and should be used for when the mouse
-   * cursor exits the link toolbar.
-   */
-  public stopHideTimer = () => {
-    this.view!.stopMenuUpdateTimer();
-  };
-
-  public get shown() {
-    return this.view?.state?.show || false;
-  }
-
-  public closeMenu = () => this.view!.closeMenu();
-}
diff --git a/packages/core/src/extensions/NodeSelectionKeyboard/NodeSelectionKeyboard.ts b/packages/core/src/extensions/NodeSelectionKeyboard/NodeSelectionKeyboard.ts
new file mode 100644
index 0000000000..b5d7f9394b
--- /dev/null
+++ b/packages/core/src/extensions/NodeSelectionKeyboard/NodeSelectionKeyboard.ts
@@ -0,0 +1,74 @@
+import { Plugin, PluginKey, TextSelection } from "prosemirror-state";
+import { createExtension } from "../../editor/BlockNoteExtension.js";
+
+const PLUGIN_KEY = new PluginKey("node-selection-keyboard");
+// By default, typing with a node selection active will cause ProseMirror to
+// replace the node with one that contains editable content. This plugin blocks
+// this behaviour without also blocking things like keyboard shortcuts:
+//
+// - Lets through key presses that do not include alphanumeric characters. This
+// includes things like backspace/delete/home/end/etc.
+// - Lets through any key presses that include ctrl/meta keys. These will be
+// shortcuts of some kind like ctrl+C/mod+C.
+// - Special case for Enter key which creates a new paragraph block below and
+// sets the selection to it. This is just to bring the UX closer to Notion
+//
+// While a more elegant solution would probably process transactions instead of
+// keystrokes, this brings us most of the way to Notion's UX without much added
+// complexity.
+export const NodeSelectionKeyboard = createExtension(
+  () =>
+    ({
+      key: "nodeSelectionKeyboard",
+      prosemirrorPlugins: [
+        new Plugin({
+          key: PLUGIN_KEY,
+          props: {
+            handleKeyDown: (view, event) => {
+              // Checks for node selection
+              if ("node" in view.state.selection) {
+                // Checks if key press uses ctrl/meta modifier
+                if (event.ctrlKey || event.metaKey) {
+                  return false;
+                }
+                // Checks if key press is alphanumeric
+                if (event.key.length === 1) {
+                  event.preventDefault();
+
+                  return true;
+                }
+                // Checks if key press is Enter
+                if (
+                  event.key === "Enter" &&
+                  !event.shiftKey &&
+                  !event.altKey &&
+                  !event.ctrlKey &&
+                  !event.metaKey
+                ) {
+                  const tr = view.state.tr;
+                  view.dispatch(
+                    tr
+                      .insert(
+                        view.state.tr.selection.$to.after(),
+                        view.state.schema.nodes["paragraph"].createChecked(),
+                      )
+                      .setSelection(
+                        new TextSelection(
+                          tr.doc.resolve(
+                            view.state.tr.selection.$to.after() + 1,
+                          ),
+                        ),
+                      ),
+                  );
+
+                  return true;
+                }
+              }
+
+              return false;
+            },
+          },
+        }),
+      ],
+    }) as const,
+);
diff --git a/packages/core/src/extensions/NodeSelectionKeyboard/NodeSelectionKeyboardPlugin.ts b/packages/core/src/extensions/NodeSelectionKeyboard/NodeSelectionKeyboardPlugin.ts
deleted file mode 100644
index 0b1428b637..0000000000
--- a/packages/core/src/extensions/NodeSelectionKeyboard/NodeSelectionKeyboardPlugin.ts
+++ /dev/null
@@ -1,75 +0,0 @@
-import { Plugin, PluginKey, TextSelection } from "prosemirror-state";
-import { BlockNoteExtension } from "../../editor/BlockNoteExtension.js";
-
-const PLUGIN_KEY = new PluginKey("node-selection-keyboard");
-// By default, typing with a node selection active will cause ProseMirror to
-// replace the node with one that contains editable content. This plugin blocks
-// this behaviour without also blocking things like keyboard shortcuts:
-//
-// - Lets through key presses that do not include alphanumeric characters. This
-// includes things like backspace/delete/home/end/etc.
-// - Lets through any key presses that include ctrl/meta keys. These will be
-// shortcuts of some kind like ctrl+C/mod+C.
-// - Special case for Enter key which creates a new paragraph block below and
-// sets the selection to it. This is just to bring the UX closer to Notion
-//
-// While a more elegant solution would probably process transactions instead of
-// keystrokes, this brings us most of the way to Notion's UX without much added
-// complexity.
-export class NodeSelectionKeyboardPlugin extends BlockNoteExtension {
-  public static key() {
-    return "nodeSelectionKeyboard";
-  }
-
-  constructor() {
-    super();
-    this.addProsemirrorPlugin(
-      new Plugin({
-        key: PLUGIN_KEY,
-        props: {
-          handleKeyDown: (view, event) => {
-            // Checks for node selection
-            if ("node" in view.state.selection) {
-              // Checks if key press uses ctrl/meta modifier
-              if (event.ctrlKey || event.metaKey) {
-                return false;
-              }
-              // Checks if key press is alphanumeric
-              if (event.key.length === 1) {
-                event.preventDefault();
-
-                return true;
-              }
-              // Checks if key press is Enter
-              if (
-                event.key === "Enter" &&
-                !event.shiftKey &&
-                !event.altKey &&
-                !event.ctrlKey &&
-                !event.metaKey
-              ) {
-                const tr = view.state.tr;
-                view.dispatch(
-                  tr
-                    .insert(
-                      view.state.tr.selection.$to.after(),
-                      view.state.schema.nodes["paragraph"].createChecked(),
-                    )
-                    .setSelection(
-                      new TextSelection(
-                        tr.doc.resolve(view.state.tr.selection.$to.after() + 1),
-                      ),
-                    ),
-                );
-
-                return true;
-              }
-            }
-
-            return false;
-          },
-        },
-      }),
-    );
-  }
-}
diff --git a/packages/core/src/extensions/Placeholder/PlaceholderPlugin.ts b/packages/core/src/extensions/Placeholder/Placeholder.ts
similarity index 89%
rename from packages/core/src/extensions/Placeholder/PlaceholderPlugin.ts
rename to packages/core/src/extensions/Placeholder/Placeholder.ts
index b2bdc46858..367187fd04 100644
--- a/packages/core/src/extensions/Placeholder/PlaceholderPlugin.ts
+++ b/packages/core/src/extensions/Placeholder/Placeholder.ts
@@ -1,25 +1,16 @@
 import { Plugin, PluginKey } from "prosemirror-state";
 import { Decoration, DecorationSet } from "prosemirror-view";
 import { v4 } from "uuid";
-import type { BlockNoteEditor } from "../../editor/BlockNoteEditor.js";
-import { BlockNoteExtension } from "../../editor/BlockNoteExtension.js";
+import { createExtension } from "../../editor/BlockNoteExtension.js";
 
 const PLUGIN_KEY = new PluginKey(`blocknote-placeholder`);
 
-export class PlaceholderPlugin extends BlockNoteExtension {
-  public static key() {
-    return "placeholder";
-  }
-
-  constructor(
-    editor: BlockNoteEditor<any, any, any>,
-    placeholders: Record<
-      string | "default" | "emptyDocument",
-      string | undefined
-    >,
-  ) {
-    super();
-    this.addProsemirrorPlugin(
+export const Placeholder = createExtension((editor, options) => {
+  // TODO defaults?
+  const placeholders = options.placeholders;
+  return {
+    key: "placeholder",
+    prosemirrorPlugins: [
       new Plugin({
         key: PLUGIN_KEY,
         view: (view) => {
@@ -49,7 +40,7 @@ export class PlaceholderPlugin extends BlockNoteExtension {
               default: defaultPlaceholder,
               emptyDocument: emptyPlaceholder,
               ...rest
-            } = placeholders;
+            } = placeholders || {};
 
             // add block specific placeholders
             for (const [blockType, placeholder] of Object.entries(rest)) {
@@ -142,6 +133,6 @@ export class PlaceholderPlugin extends BlockNoteExtension {
           },
         },
       }),
-    );
-  }
-}
+    ],
+  } as const;
+});
diff --git a/packages/core/src/extensions/PreviousBlockType/PreviousBlockTypePlugin.ts b/packages/core/src/extensions/PreviousBlockType/PreviousBlockType.ts
similarity index 95%
rename from packages/core/src/extensions/PreviousBlockType/PreviousBlockTypePlugin.ts
rename to packages/core/src/extensions/PreviousBlockType/PreviousBlockType.ts
index 587c005aab..1ac88eb217 100644
--- a/packages/core/src/extensions/PreviousBlockType/PreviousBlockTypePlugin.ts
+++ b/packages/core/src/extensions/PreviousBlockType/PreviousBlockType.ts
@@ -1,7 +1,7 @@
 import { findChildren } from "@tiptap/core";
 import { Plugin, PluginKey } from "prosemirror-state";
 import { Decoration, DecorationSet } from "prosemirror-view";
-import { BlockNoteExtension } from "../../editor/BlockNoteExtension.js";
+import { createExtension } from "../../editor/BlockNoteExtension.js";
 
 const PLUGIN_KEY = new PluginKey(`previous-blocks`);
 
@@ -24,15 +24,14 @@ const nodeAttributes: Record<string, string> = {
  *
  * Solution: When attributes change on a node, this plugin sets a data-* attribute with the "previous" value. This way we can still use CSS transitions. (See block.module.css)
  */
-export class PreviousBlockTypePlugin extends BlockNoteExtension {
-  public static key() {
-    return "previousBlockType";
+export const PreviousBlockType = createExtension((_editor, options) => {
+  if (options.animations === false) {
+    return;
   }
-
-  constructor() {
-    super();
-    let timeout: ReturnType<typeof setTimeout>;
-    this.addProsemirrorPlugin(
+  let timeout: ReturnType<typeof setTimeout>;
+  return {
+    key: "previousBlockType",
+    prosemirrorPlugins: [
       new Plugin({
         key: PLUGIN_KEY,
         view(_editorView) {
@@ -207,6 +206,6 @@ export class PreviousBlockTypePlugin extends BlockNoteExtension {
           },
         },
       }),
-    );
-  }
-}
+    ],
+  } as const;
+});
diff --git a/packages/core/src/extensions/ShowSelection/ShowSelection.ts b/packages/core/src/extensions/ShowSelection/ShowSelection.ts
new file mode 100644
index 0000000000..d7f9e522a3
--- /dev/null
+++ b/packages/core/src/extensions/ShowSelection/ShowSelection.ts
@@ -0,0 +1,54 @@
+import { Plugin, PluginKey } from "prosemirror-state";
+import { Decoration, DecorationSet } from "prosemirror-view";
+import { BlockNoteEditor } from "../../editor/BlockNoteEditor.js";
+import {
+  createExtension,
+  createStore,
+} from "../../editor/BlockNoteExtension.js";
+
+const PLUGIN_KEY = new PluginKey(`blocknote-show-selection`);
+
+/**
+ * Plugin that shows adds a decoration around the current selection
+ * This can be used to highlight the current selection in the UI even when the
+ * text editor is not focused.
+ */
+export const ShowSelection = createExtension(
+  (editor: BlockNoteEditor<any, any, any>, _options) => {
+    const store = createStore(
+      { enabled: false },
+      {
+        onUpdate() {
+          editor.transact((tr) => tr.setMeta(PLUGIN_KEY, {}));
+        },
+      },
+    );
+    return {
+      key: "showSelection",
+      store,
+      prosemirrorPlugins: [
+        new Plugin({
+          key: PLUGIN_KEY,
+          props: {
+            decorations: (state) => {
+              const { doc, selection } = state;
+              if (!store.state.enabled) {
+                return DecorationSet.empty;
+              }
+              const dec = Decoration.inline(selection.from, selection.to, {
+                "data-show-selection": "true",
+              });
+              return DecorationSet.create(doc, [dec]);
+            },
+          },
+        }),
+      ],
+      /**
+       * Show or hide the selection decoration
+       */
+      showSelection(shouldShow: boolean) {
+        store.setState({ enabled: shouldShow });
+      },
+    } as const;
+  },
+);
diff --git a/packages/core/src/extensions/ShowSelection/ShowSelectionPlugin.ts b/packages/core/src/extensions/ShowSelection/ShowSelectionPlugin.ts
deleted file mode 100644
index 0cf7ed2716..0000000000
--- a/packages/core/src/extensions/ShowSelection/ShowSelectionPlugin.ts
+++ /dev/null
@@ -1,57 +0,0 @@
-import { Plugin, PluginKey } from "prosemirror-state";
-import { Decoration, DecorationSet } from "prosemirror-view";
-import { BlockNoteEditor } from "../../editor/BlockNoteEditor.js";
-import { BlockNoteExtension } from "../../editor/BlockNoteExtension.js";
-
-const PLUGIN_KEY = new PluginKey(`blocknote-show-selection`);
-
-/**
- * Plugin that shows adds a decoration around the current selection
- * This can be used to highlight the current selection in the UI even when the
- * text editor is not focused.
- */
-export class ShowSelectionPlugin extends BlockNoteExtension {
-  public static key() {
-    return "showSelection";
-  }
-
-  private enabled = false;
-
-  public constructor(private readonly editor: BlockNoteEditor<any, any, any>) {
-    super();
-    this.addProsemirrorPlugin(
-      new Plugin({
-        key: PLUGIN_KEY,
-        props: {
-          decorations: (state) => {
-            const { doc, selection } = state;
-
-            if (!this.enabled) {
-              return DecorationSet.empty;
-            }
-
-            const dec = Decoration.inline(selection.from, selection.to, {
-              "data-show-selection": "true",
-            });
-
-            return DecorationSet.create(doc, [dec]);
-          },
-        },
-      }),
-    );
-  }
-
-  public setEnabled(enabled: boolean) {
-    if (this.enabled === enabled) {
-      return;
-    }
-
-    this.enabled = enabled;
-
-    this.editor.transact((tr) => tr.setMeta(PLUGIN_KEY, {}));
-  }
-
-  public getEnabled() {
-    return this.enabled;
-  }
-}
diff --git a/packages/core/src/extensions/SideMenu/SideMenuPlugin.ts b/packages/core/src/extensions/SideMenu/SideMenu.ts
similarity index 92%
rename from packages/core/src/extensions/SideMenu/SideMenuPlugin.ts
rename to packages/core/src/extensions/SideMenu/SideMenu.ts
index 0ba69b9893..312b2f3662 100644
--- a/packages/core/src/extensions/SideMenu/SideMenuPlugin.ts
+++ b/packages/core/src/extensions/SideMenu/SideMenu.ts
@@ -10,7 +10,10 @@ import { EditorView } from "@tiptap/pm/view";
 
 import { Block } from "../../blocks/defaultBlocks.js";
 import type { BlockNoteEditor } from "../../editor/BlockNoteEditor.js";
-import { BlockNoteExtension } from "../../editor/BlockNoteExtension.js";
+import {
+  createExtension,
+  createStore,
+} from "../../editor/BlockNoteExtension.js";
 import { UiElementPosition } from "../../extensions-shared/UiElementPosition.js";
 import {
   BlockSchema,
@@ -682,83 +685,74 @@ export class SideMenuView<
 
 export const sideMenuPluginKey = new PluginKey("SideMenuPlugin");
 
-export class SideMenuProsemirrorPlugin<
-  BSchema extends BlockSchema,
-  I extends InlineContentSchema,
-  S extends StyleSchema,
-> extends BlockNoteExtension {
-  public static key() {
-    return "sideMenu";
-  }
-
-  public view: SideMenuView<BSchema, I, S> | undefined;
+export const SideMenu = createExtension((editor) => {
+  let view: SideMenuView<any, any, any> | undefined;
+  const store = createStore<SideMenuState<any, any, any> | undefined>(
+    undefined,
+  );
 
-  constructor(private readonly editor: BlockNoteEditor<BSchema, I, S>) {
-    super();
-    this.addProsemirrorPlugin(
+  return {
+    key: "sideMenu",
+    store,
+    prosemirrorPlugins: [
       new Plugin({
         key: sideMenuPluginKey,
         view: (editorView) => {
-          this.view = new SideMenuView(editor, editorView, (state) => {
-            this.emit("update", state);
+          view = new SideMenuView(editor, editorView, (state) => {
+            // TODO: Without spreading the state, in some cases like toggling
+            // `show`, this doesn't trigger an update.
+            store.setState({ ...state });
           });
-          return this.view;
+          return view;
         },
       }),
-    );
-  }
-
-  public onUpdate(callback: (state: SideMenuState<BSchema, I, S>) => void) {
-    return this.on("update", callback);
-  }
-
-  /**
-   * Handles drag & drop events for blocks.
-   */
-  blockDragStart = (
-    event: {
-      dataTransfer: DataTransfer | null;
-      clientY: number;
+    ],
+
+    /**
+     * Handles drag & drop events for blocks.
+     */
+    blockDragStart(
+      event: { dataTransfer: DataTransfer | null; clientY: number },
+      block: Block<any, any, any>,
+    ) {
+      if (view) {
+        view.isDragOrigin = true;
+      }
+      dragStart(event, block, editor);
     },
-    block: Block<BSchema, I, S>,
-  ) => {
-    if (this.view) {
-      this.view.isDragOrigin = true;
-    }
 
-    dragStart(event, block, this.editor);
-  };
+    /**
+     * Handles drag & drop events for blocks.
+     */
+    blockDragEnd() {
+      unsetDragImage(editor.prosemirrorView.root);
+      if (view) {
+        view.isDragOrigin = false;
+      }
 
-  /**
-   * Handles drag & drop events for blocks.
-   */
-  blockDragEnd = () => {
-    unsetDragImage(this.editor.prosemirrorView.root);
+      editor.blur();
+    },
 
-    if (this.view) {
-      this.view.isDragOrigin = false;
-    }
+    /**
+     * Freezes the side menu. When frozen, the side menu will stay
+     * attached to the same block regardless of which block is hovered by the
+     * mouse cursor.
+     */
+    freezeMenu() {
+      view!.menuFrozen = true;
+      view!.state!.show = true;
+      view!.emitUpdate(view!.state!);
+    },
 
-    this.editor.blur();
-  };
-  /**
-   * Freezes the side menu. When frozen, the side menu will stay
-   * attached to the same block regardless of which block is hovered by the
-   * mouse cursor.
-   */
-  freezeMenu = () => {
-    this.view!.menuFrozen = true;
-    this.view!.state!.show = true;
-    this.view!.emitUpdate(this.view!.state!);
-  };
-  /**
-   * Unfreezes the side menu. When frozen, the side menu will stay
-   * attached to the same block regardless of which block is hovered by the
-   * mouse cursor.
-   */
-  unfreezeMenu = () => {
-    this.view!.menuFrozen = false;
-    this.view!.state!.show = false;
-    this.view!.emitUpdate(this.view!.state!);
-  };
-}
+    /**
+     * Unfreezes the side menu. When frozen, the side menu will stay
+     * attached to the same block regardless of which block is hovered by the
+     * mouse cursor.
+     */
+    unfreezeMenu() {
+      view!.menuFrozen = false;
+      view!.state!.show = false;
+      view!.emitUpdate(view!.state!);
+    },
+  } as const;
+});
diff --git a/packages/core/src/extensions/SuggestionMenu/SuggestionPlugin.ts b/packages/core/src/extensions/SuggestionMenu/SuggestionMenu.ts
similarity index 82%
rename from packages/core/src/extensions/SuggestionMenu/SuggestionPlugin.ts
rename to packages/core/src/extensions/SuggestionMenu/SuggestionMenu.ts
index 0c4bb7f4a6..a8be8d1b1b 100644
--- a/packages/core/src/extensions/SuggestionMenu/SuggestionPlugin.ts
+++ b/packages/core/src/extensions/SuggestionMenu/SuggestionMenu.ts
@@ -3,14 +3,12 @@ import { EditorState, Plugin, PluginKey } from "prosemirror-state";
 import { Decoration, DecorationSet, EditorView } from "prosemirror-view";
 
 import { trackPosition } from "../../api/positionMapping.js";
-import type { BlockNoteEditor } from "../../editor/BlockNoteEditor.js";
-import { BlockNoteExtension } from "../../editor/BlockNoteExtension.js";
-import { UiElementPosition } from "../../extensions-shared/UiElementPosition.js";
 import {
-  BlockSchema,
-  InlineContentSchema,
-  StyleSchema,
-} from "../../schema/index.js";
+  createExtension,
+  createStore,
+} from "../../editor/BlockNoteExtension.js";
+import { UiElementPosition } from "../../extensions-shared/UiElementPosition.js";
+import { BlockNoteEditor } from "../../editor/BlockNoteEditor.js";
 
 const findBlock = findParentNode((node) => node.type.name === "blockContainer");
 
@@ -19,18 +17,14 @@ export type SuggestionMenuState = UiElementPosition & {
   ignoreQueryLength?: boolean;
 };
 
-class SuggestionMenuView<
-  BSchema extends BlockSchema,
-  I extends InlineContentSchema,
-  S extends StyleSchema,
-> {
+class SuggestionMenuView {
   public state?: SuggestionMenuState;
   public emitUpdate: (triggerCharacter: string) => void;
   private rootEl?: Document | ShadowRoot;
   pluginState: SuggestionPluginState;
 
   constructor(
-    private readonly editor: BlockNoteEditor<BSchema, I, S>,
+    private readonly editor: BlockNoteEditor<any, any, any>,
     emitUpdate: (menuName: string, state: SuggestionMenuState) => void,
     view: EditorView,
   ) {
@@ -63,7 +57,9 @@ class SuggestionMenuView<
       if (!decorationNode) {
         return;
       }
-      this.state.referencePos = decorationNode.getBoundingClientRect();
+      this.state.referencePos = decorationNode
+        .getBoundingClientRect()
+        .toJSON() as DOMRect;
       this.emitUpdate(this.pluginState!.triggerCharacter!);
     }
   };
@@ -103,7 +99,9 @@ class SuggestionMenuView<
     if (this.editor.isEditable && decorationNode) {
       this.state = {
         show: true,
-        referencePos: decorationNode.getBoundingClientRect(),
+        referencePos: decorationNode
+          .getBoundingClientRect()
+          .toJSON() as DOMRect,
         query: this.pluginState!.query,
       };
 
@@ -163,34 +161,68 @@ const suggestionMenuPluginKey = new PluginKey("SuggestionMenuPlugin");
  * - This version hides some unnecessary complexity from the user of the plugin.
  * - This version handles key events differently
  */
-export class SuggestionMenuProseMirrorPlugin<
-  BSchema extends BlockSchema,
-  I extends InlineContentSchema,
-  S extends StyleSchema,
-> extends BlockNoteExtension {
-  public static key() {
-    return "suggestionMenu";
-  }
-
-  private view: SuggestionMenuView<BSchema, I, S> | undefined;
-  private triggerCharacters: string[] = [];
+export const SuggestionMenu = createExtension((editor) => {
+  const triggerCharacters: string[] = [];
+  let view: SuggestionMenuView | undefined = undefined;
+  const store = createStore<
+    (SuggestionMenuState & { triggerCharacter: string }) | undefined
+  >(undefined);
+  return {
+    key: "suggestionMenu",
+    store,
+    addTriggerCharacter: (triggerCharacter: string) => {
+      triggerCharacters.push(triggerCharacter);
+    },
+    removeTriggerCharacter: (triggerCharacter: string) => {
+      triggerCharacters.splice(triggerCharacters.indexOf(triggerCharacter), 1);
+    },
+    closeMenu: () => {
+      view?.closeMenu();
+    },
+    clearQuery: () => {
+      view?.clearQuery();
+    },
+    shown: () => {
+      return view?.state?.show || false;
+    },
+    openSuggestionMenu: (
+      triggerCharacter: string,
+      pluginState?: {
+        deleteTriggerCharacter?: boolean;
+        ignoreQueryLength?: boolean;
+      },
+    ) => {
+      if (editor.headless) {
+        return;
+      }
 
-  constructor(editor: BlockNoteEditor<BSchema, I, S>) {
-    super();
-    const triggerCharacters = this.triggerCharacters;
-    this.addProsemirrorPlugin(
+      editor.focus();
+
+      editor.transact((tr) => {
+        if (pluginState?.deleteTriggerCharacter) {
+          tr.insertText(triggerCharacter);
+        }
+        tr.scrollIntoView().setMeta(suggestionMenuPluginKey, {
+          triggerCharacter: triggerCharacter,
+          deleteTriggerCharacter: pluginState?.deleteTriggerCharacter || false,
+          ignoreQueryLength: pluginState?.ignoreQueryLength || false,
+        });
+      });
+    },
+    // TODO this whole plugin needs to be refactored (but I've done the minimal)
+    prosemirrorPlugins: [
       new Plugin({
         key: suggestionMenuPluginKey,
 
-        view: (view) => {
-          this.view = new SuggestionMenuView<BSchema, I, S>(
+        view: (v) => {
+          view = new SuggestionMenuView(
             editor,
             (triggerCharacter, state) => {
-              this.emit(`update ${triggerCharacter}`, state);
+              store.setState({ ...state, triggerCharacter });
             },
-            view,
+            v,
           );
-          return this.view;
+          return view;
         },
 
         state: {
@@ -225,7 +257,7 @@ export class SuggestionMenuProseMirrorPlugin<
             ) {
               if (prev) {
                 // Close the previous menu if it exists
-                this.closeMenu();
+                view?.closeMenu();
               }
               const trackedPosition = trackPosition(
                 editor,
@@ -360,44 +392,6 @@ export class SuggestionMenuProseMirrorPlugin<
           },
         },
       }),
-    );
-  }
-
-  public onUpdate(
-    triggerCharacter: string,
-    callback: (state: SuggestionMenuState) => void,
-  ) {
-    if (!this.triggerCharacters.includes(triggerCharacter)) {
-      this.addTriggerCharacter(triggerCharacter);
-    }
-    // TODO: be able to remove the triggerCharacter
-    return this.on(`update ${triggerCharacter}`, callback);
-  }
-
-  addTriggerCharacter = (triggerCharacter: string) => {
-    this.triggerCharacters.push(triggerCharacter);
-  };
-
-  // TODO: Should this be called automatically when listeners are removed?
-  removeTriggerCharacter = (triggerCharacter: string) => {
-    this.triggerCharacters = this.triggerCharacters.filter(
-      (c) => c !== triggerCharacter,
-    );
-  };
-
-  closeMenu = () => this.view!.closeMenu();
-
-  clearQuery = () => this.view!.clearQuery();
-
-  public get shown() {
-    return this.view?.state?.show || false;
-  }
-}
-
-export function createSuggestionMenu<
-  BSchema extends BlockSchema,
-  I extends InlineContentSchema,
-  S extends StyleSchema,
->(editor: BlockNoteEditor<BSchema, I, S>, triggerCharacter: string) {
-  editor.suggestionMenus.addTriggerCharacter(triggerCharacter);
-}
+    ],
+  } as const;
+});
diff --git a/packages/core/src/extensions/SuggestionMenu/getDefaultSlashMenuItems.ts b/packages/core/src/extensions/SuggestionMenu/getDefaultSlashMenuItems.ts
index 921d181d1a..2b50c795ba 100644
--- a/packages/core/src/extensions/SuggestionMenu/getDefaultSlashMenuItems.ts
+++ b/packages/core/src/extensions/SuggestionMenu/getDefaultSlashMenuItems.ts
@@ -10,6 +10,8 @@ import {
 } from "../../schema/index.js";
 import { formatKeyboardShortcut } from "../../util/browser.js";
 import { DefaultSuggestionItem } from "./DefaultSuggestionItem.js";
+import { FilePanel } from "../FilePanel/FilePanel.js";
+import { SuggestionMenu } from "./SuggestionMenu.js";
 
 // Sets the editor's text cursor position to the next content editable block,
 // so either a block with inline content or a table. The last block is always a
@@ -41,7 +43,7 @@ function setSelectionToNextContentEditableBlock<
 // updates the current block instead of inserting a new one below. If the new
 // block doesn't contain editable content, the cursor is moved to the next block
 // that does.
-export function insertOrUpdateBlock<
+export function insertOrUpdateBlockForSlashMenu<
   BSchema extends BlockSchema,
   I extends InlineContentSchema,
   S extends StyleSchema,
@@ -91,7 +93,7 @@ export function getDefaultSlashMenuItems<
     items.push(
       {
         onItemClick: () => {
-          insertOrUpdateBlock(editor, {
+          insertOrUpdateBlockForSlashMenu(editor, {
             type: "heading",
             props: { level: 1 },
           });
@@ -102,7 +104,7 @@ export function getDefaultSlashMenuItems<
       },
       {
         onItemClick: () => {
-          insertOrUpdateBlock(editor, {
+          insertOrUpdateBlockForSlashMenu(editor, {
             type: "heading",
             props: { level: 2 },
           });
@@ -113,7 +115,7 @@ export function getDefaultSlashMenuItems<
       },
       {
         onItemClick: () => {
-          insertOrUpdateBlock(editor, {
+          insertOrUpdateBlockForSlashMenu(editor, {
             type: "heading",
             props: { level: 3 },
           });
@@ -128,7 +130,7 @@ export function getDefaultSlashMenuItems<
   if (editorHasBlockWithType(editor, "quote")) {
     items.push({
       onItemClick: () => {
-        insertOrUpdateBlock(editor, {
+        insertOrUpdateBlockForSlashMenu(editor, {
           type: "quote",
         });
       },
@@ -140,7 +142,7 @@ export function getDefaultSlashMenuItems<
   if (editorHasBlockWithType(editor, "toggleListItem")) {
     items.push({
       onItemClick: () => {
-        insertOrUpdateBlock(editor, {
+        insertOrUpdateBlockForSlashMenu(editor, {
           type: "toggleListItem",
         });
       },
@@ -153,7 +155,7 @@ export function getDefaultSlashMenuItems<
   if (editorHasBlockWithType(editor, "numberedListItem")) {
     items.push({
       onItemClick: () => {
-        insertOrUpdateBlock(editor, {
+        insertOrUpdateBlockForSlashMenu(editor, {
           type: "numberedListItem",
         });
       },
@@ -166,7 +168,7 @@ export function getDefaultSlashMenuItems<
   if (editorHasBlockWithType(editor, "bulletListItem")) {
     items.push({
       onItemClick: () => {
-        insertOrUpdateBlock(editor, {
+        insertOrUpdateBlockForSlashMenu(editor, {
           type: "bulletListItem",
         });
       },
@@ -179,7 +181,7 @@ export function getDefaultSlashMenuItems<
   if (editorHasBlockWithType(editor, "checkListItem")) {
     items.push({
       onItemClick: () => {
-        insertOrUpdateBlock(editor, {
+        insertOrUpdateBlockForSlashMenu(editor, {
           type: "checkListItem",
         });
       },
@@ -192,7 +194,7 @@ export function getDefaultSlashMenuItems<
   if (editorHasBlockWithType(editor, "paragraph")) {
     items.push({
       onItemClick: () => {
-        insertOrUpdateBlock(editor, {
+        insertOrUpdateBlockForSlashMenu(editor, {
           type: "paragraph",
         });
       },
@@ -205,7 +207,7 @@ export function getDefaultSlashMenuItems<
   if (editorHasBlockWithType(editor, "codeBlock")) {
     items.push({
       onItemClick: () => {
-        insertOrUpdateBlock(editor, {
+        insertOrUpdateBlockForSlashMenu(editor, {
           type: "codeBlock",
         });
       },
@@ -218,7 +220,7 @@ export function getDefaultSlashMenuItems<
   if (editorHasBlockWithType(editor, "divider")) {
     items.push({
       onItemClick: () => {
-        insertOrUpdateBlock(editor, { type: "divider" });
+        insertOrUpdateBlockForSlashMenu(editor, { type: "divider" });
       },
       key: "divider",
       ...editor.dictionary.slash_menu.divider,
@@ -228,7 +230,7 @@ export function getDefaultSlashMenuItems<
   if (editorHasBlockWithType(editor, "table")) {
     items.push({
       onItemClick: () => {
-        insertOrUpdateBlock(editor, {
+        insertOrUpdateBlockForSlashMenu(editor, {
           type: "table",
           content: {
             type: "tableContent",
@@ -252,16 +254,12 @@ export function getDefaultSlashMenuItems<
   if (editorHasBlockWithType(editor, "image", { url: "string" })) {
     items.push({
       onItemClick: () => {
-        const insertedBlock = insertOrUpdateBlock(editor, {
+        const insertedBlock = insertOrUpdateBlockForSlashMenu(editor, {
           type: "image",
         });
 
         // Immediately open the file toolbar
-        editor.transact((tr) =>
-          tr.setMeta(editor.filePanel!.plugins[0], {
-            block: insertedBlock,
-          }),
-        );
+        editor.getExtension(FilePanel)?.showMenu(insertedBlock.id);
       },
       key: "image",
       ...editor.dictionary.slash_menu.image,
@@ -271,16 +269,12 @@ export function getDefaultSlashMenuItems<
   if (editorHasBlockWithType(editor, "video", { url: "string" })) {
     items.push({
       onItemClick: () => {
-        const insertedBlock = insertOrUpdateBlock(editor, {
+        const insertedBlock = insertOrUpdateBlockForSlashMenu(editor, {
           type: "video",
         });
 
         // Immediately open the file toolbar
-        editor.transact((tr) =>
-          tr.setMeta(editor.filePanel!.plugins[0], {
-            block: insertedBlock,
-          }),
-        );
+        editor.getExtension(FilePanel)?.showMenu(insertedBlock.id);
       },
       key: "video",
       ...editor.dictionary.slash_menu.video,
@@ -290,16 +284,12 @@ export function getDefaultSlashMenuItems<
   if (editorHasBlockWithType(editor, "audio", { url: "string" })) {
     items.push({
       onItemClick: () => {
-        const insertedBlock = insertOrUpdateBlock(editor, {
+        const insertedBlock = insertOrUpdateBlockForSlashMenu(editor, {
           type: "audio",
         });
 
         // Immediately open the file toolbar
-        editor.transact((tr) =>
-          tr.setMeta(editor.filePanel!.plugins[0], {
-            block: insertedBlock,
-          }),
-        );
+        editor.getExtension(FilePanel)?.showMenu(insertedBlock.id);
       },
       key: "audio",
       ...editor.dictionary.slash_menu.audio,
@@ -309,16 +299,12 @@ export function getDefaultSlashMenuItems<
   if (editorHasBlockWithType(editor, "file", { url: "string" })) {
     items.push({
       onItemClick: () => {
-        const insertedBlock = insertOrUpdateBlock(editor, {
+        const insertedBlock = insertOrUpdateBlockForSlashMenu(editor, {
           type: "file",
         });
 
         // Immediately open the file toolbar
-        editor.transact((tr) =>
-          tr.setMeta(editor.filePanel!.plugins[0], {
-            block: insertedBlock,
-          }),
-        );
+        editor.getExtension(FilePanel)?.showMenu(insertedBlock.id);
       },
       key: "file",
       ...editor.dictionary.slash_menu.file,
@@ -334,7 +320,7 @@ export function getDefaultSlashMenuItems<
     items.push(
       {
         onItemClick: () => {
-          insertOrUpdateBlock(editor, {
+          insertOrUpdateBlockForSlashMenu(editor, {
             type: "heading",
             props: { level: 1, isToggleable: true },
           });
@@ -344,7 +330,7 @@ export function getDefaultSlashMenuItems<
       },
       {
         onItemClick: () => {
-          insertOrUpdateBlock(editor, {
+          insertOrUpdateBlockForSlashMenu(editor, {
             type: "heading",
             props: { level: 2, isToggleable: true },
           });
@@ -355,7 +341,7 @@ export function getDefaultSlashMenuItems<
       },
       {
         onItemClick: () => {
-          insertOrUpdateBlock(editor, {
+          insertOrUpdateBlockForSlashMenu(editor, {
             type: "heading",
             props: { level: 3, isToggleable: true },
           });
@@ -372,7 +358,7 @@ export function getDefaultSlashMenuItems<
       .forEach((level) => {
         items.push({
           onItemClick: () => {
-            insertOrUpdateBlock(editor, {
+            insertOrUpdateBlockForSlashMenu(editor, {
               type: "heading",
               props: { level: level },
             });
@@ -385,7 +371,7 @@ export function getDefaultSlashMenuItems<
 
   items.push({
     onItemClick: () => {
-      editor.openSuggestionMenu(":", {
+      editor.getExtension(SuggestionMenu)?.openSuggestionMenu(":", {
         deleteTriggerCharacter: true,
         ignoreQueryLength: true,
       });
diff --git a/packages/core/src/extensions/TableHandles/TableHandlesPlugin.ts b/packages/core/src/extensions/TableHandles/TableHandles.ts
similarity index 63%
rename from packages/core/src/extensions/TableHandles/TableHandlesPlugin.ts
rename to packages/core/src/extensions/TableHandles/TableHandles.ts
index f90fe9e95b..36196b481d 100644
--- a/packages/core/src/extensions/TableHandles/TableHandlesPlugin.ts
+++ b/packages/core/src/extensions/TableHandles/TableHandles.ts
@@ -32,29 +32,27 @@ import {
 } from "../../blocks/defaultBlockTypeGuards.js";
 import { DefaultBlockSchema } from "../../blocks/defaultBlocks.js";
 import type { BlockNoteEditor } from "../../editor/BlockNoteEditor.js";
-import { BlockNoteExtension } from "../../editor/BlockNoteExtension.js";
+import {
+  createExtension,
+  createStore,
+} from "../../editor/BlockNoteExtension.js";
 import {
   BlockFromConfigNoChildren,
   BlockSchemaWithBlock,
-  InlineContentSchema,
-  StyleSchema,
 } from "../../schema/index.js";
 import { getDraggableBlockFromElement } from "../getDraggableBlockFromElement.js";
 
 let dragImageElement: HTMLElement | undefined;
 
 // TODO consider switching this to jotai, it is a bit messy and noisy
-export type TableHandlesState<
-  I extends InlineContentSchema,
-  S extends StyleSchema,
-> = {
+export type TableHandlesState = {
   show: boolean;
   showAddOrRemoveRowsButton: boolean;
   showAddOrRemoveColumnsButton: boolean;
   referencePosCell: DOMRect | undefined;
   referencePosTable: DOMRect;
 
-  block: BlockFromConfigNoChildren<DefaultBlockSchema["table"], I, S>;
+  block: BlockFromConfigNoChildren<DefaultBlockSchema["table"], any, any>;
   colIndex: number | undefined;
   rowIndex: number | undefined;
 
@@ -144,12 +142,8 @@ function hideElements(selector: string, rootEl: Document | ShadowRoot) {
   }
 }
 
-export class TableHandlesView<
-  I extends InlineContentSchema,
-  S extends StyleSchema,
-> implements PluginView
-{
-  public state?: TableHandlesState<I, S>;
+export class TableHandlesView implements PluginView {
+  public state?: TableHandlesState;
   public emitUpdate: () => void;
 
   public tableId: string | undefined;
@@ -165,11 +159,11 @@ export class TableHandlesView<
   constructor(
     private readonly editor: BlockNoteEditor<
       BlockSchemaWithBlock<"table", DefaultBlockSchema["table"]>,
-      I,
-      S
+      any,
+      any
     >,
     private readonly pmView: EditorView,
-    emitUpdate: (state: TableHandlesState<I, S>) => void,
+    emitUpdate: (state: TableHandlesState) => void,
   ) {
     this.emitUpdate = () => {
       if (!this.state) {
@@ -260,7 +254,7 @@ export class TableHandlesView<
     this.tableElement = blockEl.node;
 
     let tableBlock:
-      | BlockFromConfigNoChildren<DefaultBlockSchema["table"], I, S>
+      | BlockFromConfigNoChildren<DefaultBlockSchema["table"], any, any>
       | undefined;
 
     const pmNodeInfo = this.editor.transact((tr) =>
@@ -614,57 +608,52 @@ export class TableHandlesView<
 
 export const tableHandlesPluginKey = new PluginKey("TableHandlesPlugin");
 
-export class TableHandlesProsemirrorPlugin<
-  I extends InlineContentSchema,
-  S extends StyleSchema,
-> extends BlockNoteExtension {
-  public static key() {
-    return "tableHandles";
-  }
+export const TableHandles = createExtension((editor) => {
+  let view: TableHandlesView | undefined = undefined;
 
-  private view: TableHandlesView<I, S> | undefined;
+  const store = createStore<TableHandlesState | undefined>(undefined);
 
-  constructor(
-    private readonly editor: BlockNoteEditor<
-      BlockSchemaWithBlock<"table", DefaultBlockSchema["table"]>,
-      I,
-      S
-    >,
-  ) {
-    super();
-    this.addProsemirrorPlugin(
+  return {
+    key: "tableHandles",
+    store,
+    prosemirrorPlugins: [
       new Plugin({
         key: tableHandlesPluginKey,
         view: (editorView) => {
-          this.view = new TableHandlesView(editor, editorView, (state) => {
-            this.emit("update", state);
+          view = new TableHandlesView(editor as any, editorView, (state) => {
+            store.setState({
+              ...state,
+              draggingState: state.draggingState
+                ? { ...state.draggingState }
+                : undefined,
+            });
           });
-          return this.view;
+          return view;
         },
         // We use decorations to render the drop cursor when dragging a table row
         // or column. The decorations are updated in the `dragOverHandler` method.
         props: {
           decorations: (state) => {
             if (
-              this.view === undefined ||
-              this.view.state === undefined ||
-              this.view.state.draggingState === undefined ||
-              this.view.tablePos === undefined
+              view === undefined ||
+              view.state === undefined ||
+              view.state.draggingState === undefined ||
+              view.tablePos === undefined
             ) {
               return;
             }
 
             const newIndex =
-              this.view.state.draggingState.draggedCellOrientation === "row"
-                ? this.view.state.rowIndex
-                : this.view.state.colIndex;
+              view.state.draggingState.draggedCellOrientation === "row"
+                ? view.state.rowIndex
+                : view.state.colIndex;
 
             if (newIndex === undefined) {
               return;
             }
 
             const decorations: Decoration[] = [];
-            const { block, draggingState } = this.view.state;
+            const { block, draggingState } = view.state;
             const { originalIndex, draggedCellOrientation } = draggingState;
 
             // Return empty decorations if:
@@ -684,13 +673,11 @@ export class TableHandlesProsemirrorPlugin<
             }
 
             // Gets the table to show the drop cursor in.
-            const tableResolvedPos = state.doc.resolve(this.view.tablePos + 1);
+            const tableResolvedPos = state.doc.resolve(view.tablePos + 1);
 
-            if (
-              this.view.state.draggingState.draggedCellOrientation === "row"
-            ) {
+            if (view.state.draggingState.draggedCellOrientation === "row") {
               const cellsInRow = getCellsAtRowHandle(
-                this.view.state.block,
+                view.state.block,
                 newIndex,
               );
 
@@ -736,7 +723,7 @@ export class TableHandlesProsemirrorPlugin<
               });
             } else {
               const cellsInColumn = getCellsAtColumnHandle(
-                this.view.state.block,
+                view.state.block,
                 newIndex,
               );
 
@@ -788,423 +775,415 @@ export class TableHandlesProsemirrorPlugin<
           },
         },
       }),
-    );
-  }
-
-  public onUpdate(callback: (state: TableHandlesState<I, S>) => void) {
-    return this.on("update", callback);
-  }
-
-  /**
-   * Callback that should be set on the `dragStart` event for whichever element
-   * is used as the column drag handle.
-   */
-  colDragStart = (event: {
-    dataTransfer: DataTransfer | null;
-    clientX: number;
-  }) => {
-    if (
-      this.view!.state === undefined ||
-      this.view!.state.colIndex === undefined
-    ) {
-      throw new Error(
-        "Attempted to drag table column, but no table block was hovered prior.",
-      );
-    }
-
-    this.view!.state.draggingState = {
-      draggedCellOrientation: "col",
-      originalIndex: this.view!.state.colIndex,
-      mousePos: event.clientX,
-    };
-    this.view!.emitUpdate();
-
-    this.editor.transact((tr) =>
-      tr.setMeta(tableHandlesPluginKey, {
-        draggedCellOrientation:
-          this.view!.state!.draggingState!.draggedCellOrientation,
-        originalIndex: this.view!.state!.colIndex,
-        newIndex: this.view!.state!.colIndex,
-        tablePos: this.view!.tablePos,
-      }),
-    );
-
-    if (this.editor.headless) {
-      return;
-    }
-
-    setHiddenDragImage(this.editor.prosemirrorView.root);
-    event.dataTransfer!.setDragImage(dragImageElement!, 0, 0);
-    event.dataTransfer!.effectAllowed = "move";
-  };
+    ],
+
+    /**
+     * Callback that should be set on the `dragStart` event for whichever element
+     * is used as the column drag handle.
+     */
+    colDragStart(event: {
+      dataTransfer: DataTransfer | null;
+      clientX: number;
+    }) {
+      if (
+        view === undefined ||
+        view.state === undefined ||
+        view.state.colIndex === undefined
+      ) {
+        throw new Error(
+          "Attempted to drag table column, but no table block was hovered prior.",
+        );
+      }
 
-  /**
-   * Callback that should be set on the `dragStart` event for whichever element
-   * is used as the row drag handle.
-   */
-  rowDragStart = (event: {
-    dataTransfer: DataTransfer | null;
-    clientY: number;
-  }) => {
-    if (
-      this.view!.state === undefined ||
-      this.view!.state.rowIndex === undefined
-    ) {
-      throw new Error(
-        "Attempted to drag table row, but no table block was hovered prior.",
+      view.state.draggingState = {
+        draggedCellOrientation: "col",
+        originalIndex: view.state.colIndex,
+        mousePos: event.clientX,
+      };
+      view.emitUpdate();
+
+      editor.transact((tr) =>
+        tr.setMeta(tableHandlesPluginKey, {
+          draggedCellOrientation:
+            view!.state!.draggingState!.draggedCellOrientation,
+          originalIndex: view!.state!.colIndex,
+          newIndex: view!.state!.colIndex,
+          tablePos: view!.tablePos,
+        }),
       );
-    }
-
-    this.view!.state.draggingState = {
-      draggedCellOrientation: "row",
-      originalIndex: this.view!.state.rowIndex,
-      mousePos: event.clientY,
-    };
-    this.view!.emitUpdate();
-
-    this.editor.transact((tr) =>
-      tr.setMeta(tableHandlesPluginKey, {
-        draggedCellOrientation:
-          this.view!.state!.draggingState!.draggedCellOrientation,
-        originalIndex: this.view!.state!.rowIndex,
-        newIndex: this.view!.state!.rowIndex,
-        tablePos: this.view!.tablePos,
-      }),
-    );
 
-    if (this.editor.headless) {
-      return;
-    }
+      if (editor.headless) {
+        return;
+      }
 
-    setHiddenDragImage(this.editor.prosemirrorView.root);
-    event.dataTransfer!.setDragImage(dragImageElement!, 0, 0);
-    event.dataTransfer!.effectAllowed = "copyMove";
-  };
+      setHiddenDragImage(editor.prosemirrorView.root);
+      event.dataTransfer!.setDragImage(dragImageElement!, 0, 0);
+      event.dataTransfer!.effectAllowed = "move";
+    },
+
+    /**
+     * Callback that should be set on the `dragStart` event for whichever element
+     * is used as the row drag handle.
+     */
+    rowDragStart(event: {
+      dataTransfer: DataTransfer | null;
+      clientY: number;
+    }) {
+      if (view!.state === undefined || view!.state.rowIndex === undefined) {
+        throw new Error(
+          "Attempted to drag table row, but no table block was hovered prior.",
+        );
+      }
 
-  /**
-   * Callback that should be set on the `dragEnd` event for both the element
-   * used as the row drag handle, and the one used as the column drag handle.
-   */
-  dragEnd = () => {
-    if (this.view!.state === undefined) {
-      throw new Error(
-        "Attempted to drag table row, but no table block was hovered prior.",
+      view!.state.draggingState = {
+        draggedCellOrientation: "row",
+        originalIndex: view!.state.rowIndex,
+        mousePos: event.clientY,
+      };
+      view!.emitUpdate();
+
+      editor.transact((tr) =>
+        tr.setMeta(tableHandlesPluginKey, {
+          draggedCellOrientation:
+            view!.state!.draggingState!.draggedCellOrientation,
+          originalIndex: view!.state!.rowIndex,
+          newIndex: view!.state!.rowIndex,
+          tablePos: view!.tablePos,
+        }),
       );
-    }
 
-    this.view!.state.draggingState = undefined;
-    this.view!.emitUpdate();
-
-    this.editor.transact((tr) => tr.setMeta(tableHandlesPluginKey, null));
-
-    if (this.editor.headless) {
-      return;
-    }
-
-    unsetHiddenDragImage(this.editor.prosemirrorView.root);
-  };
-
-  /**
-   * Freezes the drag handles. When frozen, they will stay attached to the same
-   * cell regardless of which cell is hovered by the mouse cursor.
-   */
-  freezeHandles = () => {
-    this.view!.menuFrozen = true;
-  };
-
-  /**
-   * Unfreezes the drag handles. When frozen, they will stay attached to the
-   * same cell regardless of which cell is hovered by the mouse cursor.
-   */
-  unfreezeHandles = () => {
-    this.view!.menuFrozen = false;
-  };
+      if (editor.headless) {
+        return;
+      }
 
-  getCellsAtRowHandle = (
-    block: BlockFromConfigNoChildren<DefaultBlockSchema["table"], any, any>,
-    relativeRowIndex: RelativeCellIndices["row"],
-  ) => {
-    return getCellsAtRowHandle(block, relativeRowIndex);
-  };
+      setHiddenDragImage(editor.prosemirrorView.root);
+      event.dataTransfer!.setDragImage(dragImageElement!, 0, 0);
+      event.dataTransfer!.effectAllowed = "copyMove";
+    },
+
+    /**
+     * Callback that should be set on the `dragEnd` event for both the element
+     * used as the row drag handle, and the one used as the column drag handle.
+     */
+    dragEnd() {
+      if (view!.state === undefined) {
+        throw new Error(
+          "Attempted to drag table row, but no table block was hovered prior.",
+        );
+      }
 
-  /**
-   * Get all the cells in a column of the table block.
-   */
-  getCellsAtColumnHandle = (
-    block: BlockFromConfigNoChildren<DefaultBlockSchema["table"], any, any>,
-    relativeColumnIndex: RelativeCellIndices["col"],
-  ) => {
-    return getCellsAtColumnHandle(block, relativeColumnIndex);
-  };
+      view!.state.draggingState = undefined;
+      view!.emitUpdate();
 
-  /**
-   * Sets the selection to the given cell or a range of cells.
-   * @returns The new state after the selection has been set.
-   */
-  private setCellSelection = (
-    state: EditorState,
-    relativeStartCell: RelativeCellIndices,
-    relativeEndCell: RelativeCellIndices = relativeStartCell,
-  ) => {
-    const view = this.view;
-
-    if (!view) {
-      throw new Error("Table handles view not initialized");
-    }
+      editor.transact((tr) => tr.setMeta(tableHandlesPluginKey, null));
 
-    const tableResolvedPos = state.doc.resolve(view.tablePos! + 1);
-    const startRowResolvedPos = state.doc.resolve(
-      tableResolvedPos.posAtIndex(relativeStartCell.row) + 1,
-    );
-    const startCellResolvedPos = state.doc.resolve(
-      // No need for +1, since CellSelection expects the position before the cell
-      startRowResolvedPos.posAtIndex(relativeStartCell.col),
-    );
-    const endRowResolvedPos = state.doc.resolve(
-      tableResolvedPos.posAtIndex(relativeEndCell.row) + 1,
-    );
-    const endCellResolvedPos = state.doc.resolve(
-      // No need for +1, since CellSelection expects the position before the cell
-      endRowResolvedPos.posAtIndex(relativeEndCell.col),
-    );
+      if (editor.headless) {
+        return;
+      }
 
-    // Begin a new transaction to set the selection
-    const tr = state.tr;
+      unsetHiddenDragImage(editor.prosemirrorView.root);
+    },
+
+    /**
+     * Freezes the drag handles. When frozen, they will stay attached to the same
+     * cell regardless of which cell is hovered by the mouse cursor.
+     */
+    freezeHandles() {
+      view!.menuFrozen = true;
+    },
+
+    /**
+     * Unfreezes the drag handles. When frozen, they will stay attached to the
+     * same cell regardless of which cell is hovered by the mouse cursor.
+     */
+    unfreezeHandles() {
+      view!.menuFrozen = false;
+    },
+
+    getCellsAtRowHandle(
+      block: BlockFromConfigNoChildren<DefaultBlockSchema["table"], any, any>,
+      relativeRowIndex: RelativeCellIndices["row"],
+    ) {
+      return getCellsAtRowHandle(block, relativeRowIndex);
+    },
+
+    /**
+     * Get all the cells in a column of the table block.
+     */
+    getCellsAtColumnHandle(
+      block: BlockFromConfigNoChildren<DefaultBlockSchema["table"], any, any>,
+      relativeColumnIndex: RelativeCellIndices["col"],
+    ) {
+      return getCellsAtColumnHandle(block, relativeColumnIndex);
+    },
+
+    /**
+     * Sets the selection to the given cell or a range of cells.
+     * @returns The new state after the selection has been set.
+     */
+    setCellSelection(
+      state: EditorState,
+      relativeStartCell: RelativeCellIndices,
+      relativeEndCell: RelativeCellIndices = relativeStartCell,
+    ) {
+      if (!view) {
+        throw new Error("Table handles view not initialized");
+      }
 
-    // Set the selection to the given cell or a range of cells
-    tr.setSelection(
-      new CellSelection(startCellResolvedPos, endCellResolvedPos),
-    );
+      const tableResolvedPos = state.doc.resolve(view.tablePos! + 1);
+      const startRowResolvedPos = state.doc.resolve(
+        tableResolvedPos.posAtIndex(relativeStartCell.row) + 1,
+      );
+      const startCellResolvedPos = state.doc.resolve(
+        // No need for +1, since CellSelection expects the position before the cell
+        startRowResolvedPos.posAtIndex(relativeStartCell.col),
+      );
+      const endRowResolvedPos = state.doc.resolve(
+        tableResolvedPos.posAtIndex(relativeEndCell.row) + 1,
+      );
+      const endCellResolvedPos = state.doc.resolve(
+        // No need for +1, since CellSelection expects the position before the cell
+        endRowResolvedPos.posAtIndex(relativeEndCell.col),
+      );
 
-    // Quickly apply the transaction to get the new state to update the selection before splitting the cell
-    return state.apply(tr);
-  };
+      // Begin a new transaction to set the selection
+      const tr = state.tr;
 
-  /**
-   * Adds a row or column to the table using prosemirror-table commands
-   */
-  addRowOrColumn = (
-    index: RelativeCellIndices["row"] | RelativeCellIndices["col"],
-    direction:
-      | { orientation: "row"; side: "above" | "below" }
-      | { orientation: "column"; side: "left" | "right" },
-  ) => {
-    this.editor.exec((beforeState, dispatch) => {
-      const state = this.setCellSelection(
-        beforeState,
-        direction.orientation === "row"
-          ? { row: index, col: 0 }
-          : { row: 0, col: index },
+      // Set the selection to the given cell or a range of cells
+      tr.setSelection(
+        new CellSelection(startCellResolvedPos, endCellResolvedPos),
       );
 
-      if (direction.orientation === "row") {
-        if (direction.side === "above") {
-          return addRowBefore(state, dispatch);
+      // Quickly apply the transaction to get the new state to update the selection before splitting the cell
+      return state.apply(tr);
+    },
+
+    /**
+     * Adds a row or column to the table using prosemirror-table commands
+     */
+    addRowOrColumn(
+      index: RelativeCellIndices["row"] | RelativeCellIndices["col"],
+      direction:
+        | { orientation: "row"; side: "above" | "below" }
+        | { orientation: "column"; side: "left" | "right" },
+    ) {
+      editor.exec((beforeState, dispatch) => {
+        const state = this.setCellSelection(
+          beforeState,
+          direction.orientation === "row"
+            ? { row: index, col: 0 }
+            : { row: 0, col: index },
+        );
+
+        if (direction.orientation === "row") {
+          if (direction.side === "above") {
+            return addRowBefore(state, dispatch);
+          } else {
+            return addRowAfter(state, dispatch);
+          }
         } else {
-          return addRowAfter(state, dispatch);
+          if (direction.side === "left") {
+            return addColumnBefore(state, dispatch);
+          } else {
+            return addColumnAfter(state, dispatch);
+          }
         }
+      });
+    },
+
+    /**
+     * Removes a row or column from the table using prosemirror-table commands
+     */
+    removeRowOrColumn(
+      index: RelativeCellIndices["row"] | RelativeCellIndices["col"],
+      direction: "row" | "column",
+    ) {
+      if (direction === "row") {
+        return editor.exec((beforeState, dispatch) => {
+          const state = this.setCellSelection(beforeState, {
+            row: index,
+            col: 0,
+          });
+          return deleteRow(state, dispatch);
+        });
       } else {
-        if (direction.side === "left") {
-          return addColumnBefore(state, dispatch);
-        } else {
-          return addColumnAfter(state, dispatch);
-        }
-      }
-    });
-  };
-
-  /**
-   * Removes a row or column from the table using prosemirror-table commands
-   */
-  removeRowOrColumn = (
-    index: RelativeCellIndices["row"] | RelativeCellIndices["col"],
-    direction: "row" | "column",
-  ) => {
-    if (direction === "row") {
-      return this.editor.exec((beforeState, dispatch) => {
-        const state = this.setCellSelection(beforeState, {
-          row: index,
-          col: 0,
+        return editor.exec((beforeState, dispatch) => {
+          const state = this.setCellSelection(beforeState, {
+            row: 0,
+            col: index,
+          });
+          return deleteColumn(state, dispatch);
         });
-        return deleteRow(state, dispatch);
+      }
+    },
+
+    /**
+     * Merges the cells in the table block.
+     */
+    mergeCells(cellsToMerge?: {
+      relativeStartCell: RelativeCellIndices;
+      relativeEndCell: RelativeCellIndices;
+    }) {
+      return editor.exec((beforeState, dispatch) => {
+        const state = cellsToMerge
+          ? this.setCellSelection(
+              beforeState,
+              cellsToMerge.relativeStartCell,
+              cellsToMerge.relativeEndCell,
+            )
+          : beforeState;
+
+        return mergeCells(state, dispatch);
       });
-    } else {
-      return this.editor.exec((beforeState, dispatch) => {
-        const state = this.setCellSelection(beforeState, {
-          row: 0,
-          col: index,
-        });
-        return deleteColumn(state, dispatch);
+    },
+
+    /**
+     * Splits the cell in the table block.
+     * If no cell is provided, the current cell selected will be split.
+     */
+    splitCell(relativeCellToSplit?: RelativeCellIndices) {
+      return editor.exec((beforeState, dispatch) => {
+        const state = relativeCellToSplit
+          ? this.setCellSelection(beforeState, relativeCellToSplit)
+          : beforeState;
+
+        return splitCell(state, dispatch);
       });
-    }
-  };
-
-  /**
-   * Merges the cells in the table block.
-   */
-  mergeCells = (cellsToMerge?: {
-    relativeStartCell: RelativeCellIndices;
-    relativeEndCell: RelativeCellIndices;
-  }) => {
-    return this.editor.exec((beforeState, dispatch) => {
-      const state = cellsToMerge
-        ? this.setCellSelection(
-            beforeState,
-            cellsToMerge.relativeStartCell,
-            cellsToMerge.relativeEndCell,
-          )
-        : beforeState;
-
-      return mergeCells(state, dispatch);
-    });
-  };
-
-  /**
-   * Splits the cell in the table block.
-   * If no cell is provided, the current cell selected will be split.
-   */
-  splitCell = (relativeCellToSplit?: RelativeCellIndices) => {
-    return this.editor.exec((beforeState, dispatch) => {
-      const state = relativeCellToSplit
-        ? this.setCellSelection(beforeState, relativeCellToSplit)
-        : beforeState;
-
-      return splitCell(state, dispatch);
-    });
-  };
+    },
+
+    /**
+     * Gets the start and end cells of the current cell selection.
+     * @returns The start and end cells of the current cell selection.
+     */
+    getCellSelection():
+      | undefined
+      | {
+          from: RelativeCellIndices;
+          to: RelativeCellIndices;
+          /**
+           * All of the cells that are within the selected range.
+           */
+          cells: RelativeCellIndices[];
+        } {
+      // Based on the current selection, find the table cells that are within the selected range
+
+      return editor.transact((tr) => {
+        const selection = tr.selection;
+
+        let $fromCell = selection.$from;
+        let $toCell = selection.$to;
+        if (isTableCellSelection(selection)) {
+          // When the selection is a table cell selection, we can find the
+          // from and to cells by iterating over the ranges in the selection
+          const { ranges } = selection;
+          ranges.forEach((range) => {
+            $fromCell = range.$from.min($fromCell ?? range.$from);
+            $toCell = range.$to.max($toCell ?? range.$to);
+          });
+        } else {
+          // When the selection is a normal text selection
+          // Assumes we are within a tableParagraph
+          // And find the from and to cells by resolving the positions
+          $fromCell = tr.doc.resolve(
+            selection.$from.pos - selection.$from.parentOffset - 1,
+          );
+          $toCell = tr.doc.resolve(
+            selection.$to.pos - selection.$to.parentOffset - 1,
+          );
+
+          // Opt-out when the selection is not pointing into cells
+          if ($fromCell.pos === 0 || $toCell.pos === 0) {
+            return undefined;
+          }
+        }
 
-  /**
-   * Gets the start and end cells of the current cell selection.
-   * @returns The start and end cells of the current cell selection.
-   */
-  getCellSelection = ():
-    | undefined
-    | {
-        from: RelativeCellIndices;
-        to: RelativeCellIndices;
-        /**
-         * All of the cells that are within the selected range.
-         */
-        cells: RelativeCellIndices[];
-      } => {
-    // Based on the current selection, find the table cells that are within the selected range
-
-    return this.editor.transact((tr) => {
-      const selection = tr.selection;
-
-      let $fromCell = selection.$from;
-      let $toCell = selection.$to;
-      if (isTableCellSelection(selection)) {
-        // When the selection is a table cell selection, we can find the
-        // from and to cells by iterating over the ranges in the selection
-        const { ranges } = selection;
-        ranges.forEach((range) => {
-          $fromCell = range.$from.min($fromCell ?? range.$from);
-          $toCell = range.$to.max($toCell ?? range.$to);
-        });
-      } else {
-        // When the selection is a normal text selection
-        // Assumes we are within a tableParagraph
-        // And find the from and to cells by resolving the positions
-        $fromCell = tr.doc.resolve(
-          selection.$from.pos - selection.$from.parentOffset - 1,
-        );
-        $toCell = tr.doc.resolve(
-          selection.$to.pos - selection.$to.parentOffset - 1,
+        // Find the row and table that the from and to cells are in
+        const $fromRow = tr.doc.resolve(
+          $fromCell.pos - $fromCell.parentOffset - 1,
         );
+        const $toRow = tr.doc.resolve($toCell.pos - $toCell.parentOffset - 1);
+
+        // Find the table
+        const $table = tr.doc.resolve($fromRow.pos - $fromRow.parentOffset - 1);
+
+        // Find the column and row indices of the from and to cells
+        const fromColIndex = $fromCell.index($fromRow.depth);
+        const fromRowIndex = $fromRow.index($table.depth);
+        const toColIndex = $toCell.index($toRow.depth);
+        const toRowIndex = $toRow.index($table.depth);
+
+        const cells: RelativeCellIndices[] = [];
+        for (let row = fromRowIndex; row <= toRowIndex; row++) {
+          for (let col = fromColIndex; col <= toColIndex; col++) {
+            cells.push({ row, col });
+          }
+        }
 
-        // Opt-out when the selection is not pointing into cells
-        if ($fromCell.pos === 0 || $toCell.pos === 0) {
+        return {
+          from: {
+            row: fromRowIndex,
+            col: fromColIndex,
+          },
+          to: {
+            row: toRowIndex,
+            col: toColIndex,
+          },
+          cells,
+        };
+      });
+    },
+
+    /**
+     * Gets the direction of the merge based on the current cell selection.
+     *
+     * Returns undefined when there is no cell selection, or the selection is not within a table.
+     */
+    getMergeDirection(
+      block:
+        | BlockFromConfigNoChildren<DefaultBlockSchema["table"], any, any>
+        | undefined,
+    ) {
+      return editor.transact((tr) => {
+        const isSelectingTableCells = isTableCellSelection(tr.selection)
+          ? tr.selection
+          : undefined;
+
+        if (
+          !isSelectingTableCells ||
+          !block ||
+          // Only offer the merge button if there is more than one cell selected.
+          isSelectingTableCells.ranges.length <= 1
+        ) {
           return undefined;
         }
-      }
 
-      // Find the row and table that the from and to cells are in
-      const $fromRow = tr.doc.resolve(
-        $fromCell.pos - $fromCell.parentOffset - 1,
-      );
-      const $toRow = tr.doc.resolve($toCell.pos - $toCell.parentOffset - 1);
+        const cellSelection = this.getCellSelection();
 
-      // Find the table
-      const $table = tr.doc.resolve($fromRow.pos - $fromRow.parentOffset - 1);
-
-      // Find the column and row indices of the from and to cells
-      const fromColIndex = $fromCell.index($fromRow.depth);
-      const fromRowIndex = $fromRow.index($table.depth);
-      const toColIndex = $toCell.index($toRow.depth);
-      const toRowIndex = $toRow.index($table.depth);
-
-      const cells: RelativeCellIndices[] = [];
-      for (let row = fromRowIndex; row <= toRowIndex; row++) {
-        for (let col = fromColIndex; col <= toColIndex; col++) {
-          cells.push({ row, col });
+        if (!cellSelection) {
+          return undefined;
         }
-      }
-
-      return {
-        from: {
-          row: fromRowIndex,
-          col: fromColIndex,
-        },
-        to: {
-          row: toRowIndex,
-          col: toColIndex,
-        },
-        cells,
-      };
-    });
-  };
 
-  /**
-   * Gets the direction of the merge based on the current cell selection.
-   *
-   * Returns undefined when there is no cell selection, or the selection is not within a table.
-   */
-  getMergeDirection = (
-    block:
-      | BlockFromConfigNoChildren<DefaultBlockSchema["table"], any, any>
-      | undefined,
-  ) => {
-    return this.editor.transact((tr) => {
-      const isSelectingTableCells = isTableCellSelection(tr.selection)
-        ? tr.selection
-        : undefined;
-
-      if (
-        !isSelectingTableCells ||
-        !block ||
-        // Only offer the merge button if there is more than one cell selected.
-        isSelectingTableCells.ranges.length <= 1
-      ) {
-        return undefined;
-      }
-
-      const cellSelection = this.getCellSelection();
-
-      if (!cellSelection) {
-        return undefined;
-      }
-
-      if (areInSameColumn(cellSelection.from, cellSelection.to, block)) {
-        return "vertical";
-      }
+        if (areInSameColumn(cellSelection.from, cellSelection.to, block)) {
+          return "vertical";
+        }
 
-      return "horizontal";
-    });
-  };
+        return "horizontal";
+      });
+    },
 
-  cropEmptyRowsOrColumns = (
-    block: BlockFromConfigNoChildren<DefaultBlockSchema["table"], any, any>,
-    removeEmpty: "columns" | "rows",
-  ) => {
-    return cropEmptyRowsOrColumns(block, removeEmpty);
-  };
+    cropEmptyRowsOrColumns(
+      block: BlockFromConfigNoChildren<DefaultBlockSchema["table"], any, any>,
+      removeEmpty: "columns" | "rows",
+    ) {
+      return cropEmptyRowsOrColumns(block, removeEmpty);
+    },
 
-  addRowsOrColumns = (
-    block: BlockFromConfigNoChildren<DefaultBlockSchema["table"], any, any>,
-    addType: "columns" | "rows",
-    numToAdd: number,
-  ) => {
-    return addRowsOrColumns(block, addType, numToAdd);
-  };
-}
+    addRowsOrColumns(
+      block: BlockFromConfigNoChildren<DefaultBlockSchema["table"], any, any>,
+      addType: "columns" | "rows",
+      numToAdd: number,
+    ) {
+      return addRowsOrColumns(block, addType, numToAdd);
+    },
+  } as const;
+});
diff --git a/packages/core/src/extensions/TrailingNode/TrailingNodeExtension.ts b/packages/core/src/extensions/TrailingNode/TrailingNode.ts
similarity index 83%
rename from packages/core/src/extensions/TrailingNode/TrailingNodeExtension.ts
rename to packages/core/src/extensions/TrailingNode/TrailingNode.ts
index ea70567d68..c088f5ca4a 100644
--- a/packages/core/src/extensions/TrailingNode/TrailingNodeExtension.ts
+++ b/packages/core/src/extensions/TrailingNode/TrailingNode.ts
@@ -1,5 +1,5 @@
-import { Extension } from "@tiptap/core";
 import { Plugin, PluginKey } from "prosemirror-state";
+import { createExtension } from "../../editor/BlockNoteExtension.js";
 
 // based on https://github.com/ueberdosis/tiptap/blob/40a9404c94c7fef7900610c195536384781ae101/demos/src/Experiments/TrailingNode/Vue/trailing-node.ts
 
@@ -8,24 +8,18 @@ import { Plugin, PluginKey } from "prosemirror-state";
  * - https://github.com/ueberdosis/tiptap/blob/v1/packages/tiptap-extensions/src/extensions/TrailingNode.js
  * - https://github.com/remirror/remirror/blob/e0f1bec4a1e8073ce8f5500d62193e52321155b9/packages/prosemirror-trailing-node/src/trailing-node-plugin.ts
  */
-
-export interface TrailingNodeOptions {
-  node: string;
-}
+const plugin = new PluginKey("trailingNode");
 
 /**
  * Add a trailing node to the document so the user can always click at the bottom of the document and start typing
  */
-export const TrailingNode = Extension.create<TrailingNodeOptions>({
-  name: "trailingNode",
-
-  addProseMirrorPlugins() {
-    const plugin = new PluginKey(this.name);
-    // const disabledNodes = Object.entries(this.editor.schema.nodes)
-    //   .map(([, value]) => value)
-    //   .filter((node) => this.options.notAfter.includes(node.name));
-
-    return [
+export const TrailingNode = createExtension((_editor, options) => {
+  if (options.trailingBlock === false) {
+    return;
+  }
+  return {
+    key: "trailingNode",
+    prosemirrorPlugins: [
       new Plugin({
         key: plugin,
         appendTransaction: (_, __, state) => {
@@ -80,6 +74,6 @@ export const TrailingNode = Extension.create<TrailingNodeOptions>({
           },
         },
       }),
-    ];
-  },
+    ],
+  } as const;
 });
diff --git a/packages/core/src/extensions/index.ts b/packages/core/src/extensions/index.ts
new file mode 100644
index 0000000000..0d78d7f1a7
--- /dev/null
+++ b/packages/core/src/extensions/index.ts
@@ -0,0 +1,69 @@
+import { BlockChange } from "./BlockChange/BlockChange.js";
+import { ForkYDoc } from "./Collaboration/ForkYDoc.js";
+import { SchemaMigration } from "./Collaboration/schemaMigration/SchemaMigration.js";
+import { YCursor } from "./Collaboration/YCursorPlugin.js";
+import { YSync } from "./Collaboration/YSync.js";
+import { YUndo } from "./Collaboration/YUndo.js";
+import { Comments } from "./Comments/Comments.js";
+import { DropCursor } from "./DropCursor/DropCursor.js";
+import { FilePanel } from "./FilePanel/FilePanel.js";
+import { FormattingToolbar } from "./FormattingToolbar/FormattingToolbar.js";
+import { History } from "./History/History.js";
+import { LinkToolbar } from "./LinkToolbar/LinkToolbar.js";
+import { NodeSelectionKeyboard } from "./NodeSelectionKeyboard/NodeSelectionKeyboard.js";
+import { Placeholder } from "./Placeholder/Placeholder.js";
+import { PreviousBlockType } from "./PreviousBlockType/PreviousBlockType.js";
+import { ShowSelection } from "./ShowSelection/ShowSelection.js";
+import { SideMenu } from "./SideMenu/SideMenu.js";
+import { SuggestionMenu } from "./SuggestionMenu/SuggestionMenu.js";
+import { TableHandles } from "./TableHandles/TableHandles.js";
+import { TrailingNode } from "./TrailingNode/TrailingNode.js";
+
+export * from "./BlockChange/BlockChange.js";
+export * from "./Collaboration/ForkYDoc.js";
+export * from "./Collaboration/schemaMigration/SchemaMigration.js";
+export * from "./Collaboration/YCursorPlugin.js";
+export * from "./Collaboration/YSync.js";
+export * from "./Collaboration/YUndo.js";
+export * from "./Comments/Comments.js";
+export * from "./DropCursor/DropCursor.js";
+export * from "./FilePanel/FilePanel.js";
+export * from "./FormattingToolbar/FormattingToolbar.js";
+export * from "./History/History.js";
+export * from "./LinkToolbar/LinkToolbar.js";
+export * from "./LinkToolbar/protocols.js";
+export * from "./NodeSelectionKeyboard/NodeSelectionKeyboard.js";
+export * from "./Placeholder/Placeholder.js";
+export * from "./PreviousBlockType/PreviousBlockType.js";
+export * from "./ShowSelection/ShowSelection.js";
+export * from "./SideMenu/SideMenu.js";
+export * from "./SuggestionMenu/SuggestionMenu.js";
+export * from "./SuggestionMenu/getDefaultSlashMenuItems.js";
+export * from "./SuggestionMenu/getDefaultEmojiPickerItems.js";
+export * from "./SuggestionMenu/DefaultSuggestionItem.js";
+export * from "./SuggestionMenu/DefaultGridSuggestionItem.js";
+export * from "./TableHandles/TableHandles.js";
+export * from "./TrailingNode/TrailingNode.js";
+
+export const DEFAULT_EXTENSIONS = [
+  BlockChange,
+  Comments,
+  YCursor,
+  DropCursor,
+  FilePanel,
+  ForkYDoc,
+  FormattingToolbar,
+  History,
+  LinkToolbar,
+  NodeSelectionKeyboard,
+  Placeholder,
+  PreviousBlockType,
+  SchemaMigration,
+  ShowSelection,
+  SideMenu,
+  SuggestionMenu,
+  YSync,
+  TableHandles,
+  TrailingNode,
+  YUndo,
+] as const;
diff --git a/packages/core/src/extensions/BackgroundColor/BackgroundColorExtension.ts b/packages/core/src/extensions/tiptap-extensions/BackgroundColor/BackgroundColorExtension.ts
similarity index 81%
rename from packages/core/src/extensions/BackgroundColor/BackgroundColorExtension.ts
rename to packages/core/src/extensions/tiptap-extensions/BackgroundColor/BackgroundColorExtension.ts
index 1ba499de7d..20b9641287 100644
--- a/packages/core/src/extensions/BackgroundColor/BackgroundColorExtension.ts
+++ b/packages/core/src/extensions/tiptap-extensions/BackgroundColor/BackgroundColorExtension.ts
@@ -1,5 +1,5 @@
 import { Extension } from "@tiptap/core";
-import { getBackgroundColorAttribute } from "../../blocks/defaultProps.js";
+import { getBackgroundColorAttribute } from "../../../blocks/defaultProps.js";
 
 export const BackgroundColorExtension = Extension.create({
   name: "blockBackgroundColor",
diff --git a/packages/core/src/extensions/HardBreak/HardBreak.ts b/packages/core/src/extensions/tiptap-extensions/HardBreak/HardBreak.ts
similarity index 100%
rename from packages/core/src/extensions/HardBreak/HardBreak.ts
rename to packages/core/src/extensions/tiptap-extensions/HardBreak/HardBreak.ts
diff --git a/packages/core/src/extensions/KeyboardShortcuts/KeyboardShortcutsExtension.ts b/packages/core/src/extensions/tiptap-extensions/KeyboardShortcuts/KeyboardShortcutsExtension.ts
similarity index 94%
rename from packages/core/src/extensions/KeyboardShortcuts/KeyboardShortcutsExtension.ts
rename to packages/core/src/extensions/tiptap-extensions/KeyboardShortcuts/KeyboardShortcutsExtension.ts
index da63452923..cdb4c6940a 100644
--- a/packages/core/src/extensions/KeyboardShortcuts/KeyboardShortcutsExtension.ts
+++ b/packages/core/src/extensions/tiptap-extensions/KeyboardShortcuts/KeyboardShortcutsExtension.ts
@@ -5,13 +5,13 @@ import {
   getBottomNestedBlockInfo,
   getPrevBlockInfo,
   mergeBlocksCommand,
-} from "../../api/blockManipulation/commands/mergeBlocks/mergeBlocks.js";
-import { nestBlock } from "../../api/blockManipulation/commands/nestBlock/nestBlock.js";
-import { splitBlockCommand } from "../../api/blockManipulation/commands/splitBlock/splitBlock.js";
-import { updateBlockCommand } from "../../api/blockManipulation/commands/updateBlock/updateBlock.js";
-import { getBlockInfoFromSelection } from "../../api/getBlockInfoFromPos.js";
-import { BlockNoteEditor } from "../../editor/BlockNoteEditor.js";
-import { fixColumnList } from "../../api/blockManipulation/commands/replaceBlocks/util/fixColumnList.js";
+} from "../../../api/blockManipulation/commands/mergeBlocks/mergeBlocks.js";
+import { nestBlock } from "../../../api/blockManipulation/commands/nestBlock/nestBlock.js";
+import { fixColumnList } from "../../../api/blockManipulation/commands/replaceBlocks/util/fixColumnList.js";
+import { splitBlockCommand } from "../../../api/blockManipulation/commands/splitBlock/splitBlock.js";
+import { updateBlockCommand } from "../../../api/blockManipulation/commands/updateBlock/updateBlock.js";
+import { getBlockInfoFromSelection } from "../../../api/getBlockInfoFromPos.js";
+import { BlockNoteEditor } from "../../../editor/BlockNoteEditor.js";
 
 export const KeyboardShortcutsExtension = Extension.create<{
   editor: BlockNoteEditor<any, any, any>;
@@ -489,9 +489,11 @@ export const KeyboardShortcutsExtension = Extension.create<{
       Tab: () => {
         if (
           this.options.tabBehavior !== "prefer-indent" &&
-          (this.options.editor.formattingToolbar?.shown ||
-            this.options.editor.linkToolbar?.shown ||
-            this.options.editor.filePanel?.shown)
+          true
+          // TODO what??
+          // (this.options.editor.formattingToolbar?.shown ||
+          //   this.options.editor.linkToolbar?.shown ||
+          //   this.options.editor.filePanel?.shown)
         ) {
           // don't handle tabs if a toolbar is shown, so we can tab into / out of it
           return false;
@@ -502,9 +504,11 @@ export const KeyboardShortcutsExtension = Extension.create<{
       "Shift-Tab": () => {
         if (
           this.options.tabBehavior !== "prefer-indent" &&
-          (this.options.editor.formattingToolbar?.shown ||
-            this.options.editor.linkToolbar?.shown ||
-            this.options.editor.filePanel?.shown)
+          true
+          // TODO what??
+          // (this.options.editor.formattingToolbar?.shown ||
+          //   this.options.editor.linkToolbar?.shown ||
+          //   this.options.editor.filePanel?.shown)
         ) {
           // don't handle tabs if a toolbar is shown, so we can tab into / out of it
           return false;
diff --git a/packages/core/src/extensions/Suggestions/SuggestionMarks.ts b/packages/core/src/extensions/tiptap-extensions/Suggestions/SuggestionMarks.ts
similarity index 100%
rename from packages/core/src/extensions/Suggestions/SuggestionMarks.ts
rename to packages/core/src/extensions/tiptap-extensions/Suggestions/SuggestionMarks.ts
diff --git a/packages/core/src/extensions/TextAlignment/TextAlignmentExtension.ts b/packages/core/src/extensions/tiptap-extensions/TextAlignment/TextAlignmentExtension.ts
similarity index 100%
rename from packages/core/src/extensions/TextAlignment/TextAlignmentExtension.ts
rename to packages/core/src/extensions/tiptap-extensions/TextAlignment/TextAlignmentExtension.ts
diff --git a/packages/core/src/extensions/TextColor/TextColorExtension.ts b/packages/core/src/extensions/tiptap-extensions/TextColor/TextColorExtension.ts
similarity index 81%
rename from packages/core/src/extensions/TextColor/TextColorExtension.ts
rename to packages/core/src/extensions/tiptap-extensions/TextColor/TextColorExtension.ts
index 73357e1519..a55f0a40cf 100644
--- a/packages/core/src/extensions/TextColor/TextColorExtension.ts
+++ b/packages/core/src/extensions/tiptap-extensions/TextColor/TextColorExtension.ts
@@ -1,5 +1,5 @@
 import { Extension } from "@tiptap/core";
-import { getTextColorAttribute } from "../../blocks/defaultProps.js";
+import { getTextColorAttribute } from "../../../blocks/defaultProps.js";
 
 export const TextColorExtension = Extension.create({
   name: "blockTextColor",
diff --git a/packages/core/src/extensions/UniqueID/UniqueID.ts b/packages/core/src/extensions/tiptap-extensions/UniqueID/UniqueID.ts
similarity index 100%
rename from packages/core/src/extensions/UniqueID/UniqueID.ts
rename to packages/core/src/extensions/tiptap-extensions/UniqueID/UniqueID.ts
diff --git a/packages/core/src/extensions/tiptap-extensions/index.ts b/packages/core/src/extensions/tiptap-extensions/index.ts
new file mode 100644
index 0000000000..e6fead486c
--- /dev/null
+++ b/packages/core/src/extensions/tiptap-extensions/index.ts
@@ -0,0 +1,31 @@
+import { BackgroundColorExtension } from "./BackgroundColor/BackgroundColorExtension.js";
+import { HardBreak } from "./HardBreak/HardBreak.js";
+import { KeyboardShortcutsExtension } from "./KeyboardShortcuts/KeyboardShortcutsExtension.js";
+import {
+  SuggestionAddMark,
+  SuggestionDeleteMark,
+  SuggestionModificationMark,
+} from "./Suggestions/SuggestionMarks.js";
+import { TextAlignmentExtension } from "./TextAlignment/TextAlignmentExtension.js";
+import { TextColorExtension } from "./TextColor/TextColorExtension.js";
+import { UniqueID } from "./UniqueID/UniqueID.js";
+
+export * from "./BackgroundColor/BackgroundColorExtension.js";
+export * from "./HardBreak/HardBreak.js";
+export * from "./KeyboardShortcuts/KeyboardShortcutsExtension.js";
+export * from "./Suggestions/SuggestionMarks.js";
+export * from "./TextAlignment/TextAlignmentExtension.js";
+export * from "./TextColor/TextColorExtension.js";
+export * from "./UniqueID/UniqueID.js";
+
+export const DEFAULT_TIP_TAP_EXTENSIONS = [
+  BackgroundColorExtension,
+  HardBreak,
+  KeyboardShortcutsExtension,
+  SuggestionAddMark,
+  SuggestionDeleteMark,
+  SuggestionModificationMark,
+  TextAlignmentExtension,
+  TextColorExtension,
+  UniqueID,
+];
diff --git a/packages/core/src/index.ts b/packages/core/src/index.ts
index 4b6994d5d4..07dd62bd4d 100644
--- a/packages/core/src/index.ts
+++ b/packages/core/src/index.ts
@@ -11,22 +11,10 @@ export * from "./api/pmUtil.js";
 export * from "./blocks/index.js";
 export * from "./editor/BlockNoteEditor.js";
 export * from "./editor/BlockNoteExtension.js";
-export * from "./editor/BlockNoteExtensions.js";
 export * from "./editor/defaultColors.js";
 export * from "./editor/selectionTypes.js";
 export * from "./exporter/index.js";
 export * from "./extensions-shared/UiElementPosition.js";
-export * from "./extensions/FilePanel/FilePanelPlugin.js";
-export * from "./extensions/FormattingToolbar/FormattingToolbarPlugin.js";
-export * from "./extensions/LinkToolbar/LinkToolbarPlugin.js";
-export * from "./extensions/LinkToolbar/protocols.js";
-export * from "./extensions/SideMenu/SideMenuPlugin.js";
-export * from "./extensions/SuggestionMenu/DefaultGridSuggestionItem.js";
-export * from "./extensions/SuggestionMenu/DefaultSuggestionItem.js";
-export * from "./extensions/SuggestionMenu/getDefaultEmojiPickerItems.js";
-export * from "./extensions/SuggestionMenu/getDefaultSlashMenuItems.js";
-export * from "./extensions/SuggestionMenu/SuggestionPlugin.js";
-export * from "./extensions/TableHandles/TableHandlesPlugin.js";
 export * from "./i18n/dictionary.js";
 export * from "./schema/index.js";
 export * from "./util/browser.js";
@@ -46,7 +34,7 @@ export { selectedFragmentToHTML } from "./api/clipboard/toClipboard/copyExtensio
 // Node conversions
 export * from "./api/nodeConversions/blockToNode.js";
 export * from "./api/nodeConversions/nodeToBlock.js";
-export * from "./extensions/UniqueID/UniqueID.js";
+export * from "./extensions/tiptap-extensions/UniqueID/UniqueID.js";
 
 // for server-util (TODO: maybe move):
 export * from "./api/exporters/markdown/markdownExporter.js";
diff --git a/packages/core/src/schema/blocks/createSpec.ts b/packages/core/src/schema/blocks/createSpec.ts
index 936584264f..45382528f6 100644
--- a/packages/core/src/schema/blocks/createSpec.ts
+++ b/packages/core/src/schema/blocks/createSpec.ts
@@ -2,7 +2,10 @@ import { Editor, Node } from "@tiptap/core";
 import { DOMParser, Fragment, TagParseRule } from "@tiptap/pm/model";
 import { NodeView } from "@tiptap/pm/view";
 import { mergeParagraphs } from "../../blocks/defaultBlockHelpers.js";
-import { BlockNoteExtension } from "../../editor/BlockNoteExtension.js";
+import {
+  Extension,
+  ExtensionFactory,
+} from "../../editor/BlockNoteExtension.js";
 import { PropSchema } from "../propTypes.js";
 import {
   getBlockFromPos,
@@ -131,7 +134,7 @@ export function addNodeAndExtensionsToSpec<
 >(
   blockConfig: BlockConfig<TName, TProps, TContent>,
   blockImplementation: BlockImplementation<TName, TProps, TContent>,
-  extensions?: BlockNoteExtension<any>[],
+  extensions?: (ExtensionFactory | Extension)[],
   priority?: number,
 ): LooseBlockSpec<TName, TProps, TContent> {
   const node =
@@ -296,10 +299,10 @@ export function createBlockSpec<
             options: Partial<TOptions>,
           ) => BlockImplementation<TName, TProps, TContent>),
   extensionsOrCreator?:
-    | BlockNoteExtension<any>[]
+    | ExtensionFactory[]
     | (TOptions extends undefined
-        ? () => BlockNoteExtension<any>[]
-        : (options: Partial<TOptions>) => BlockNoteExtension<any>[]),
+        ? () => ExtensionFactory[]
+        : (options: Partial<TOptions>) => ExtensionFactory[]),
 ): (options?: Partial<TOptions>) => BlockSpec<TName, TProps, TContent>;
 export function createBlockSpec<
   const TName extends string,
@@ -329,10 +332,10 @@ export function createBlockSpec<
             BlockConf["content"]
           >),
   extensionsOrCreator?:
-    | BlockNoteExtension<any>[]
+    | ExtensionFactory[]
     | (TOptions extends undefined
-        ? () => BlockNoteExtension<any>[]
-        : (options: Partial<TOptions>) => BlockNoteExtension<any>[]),
+        ? () => ExtensionFactory[]
+        : (options: Partial<TOptions>) => ExtensionFactory[]),
 ): (
   options?: Partial<TOptions>,
 ) => BlockSpec<
@@ -359,10 +362,10 @@ export function createBlockSpec<
             options: Partial<TOptions>,
           ) => BlockImplementation<TName, TProps, TContent>),
   extensionsOrCreator?:
-    | BlockNoteExtension<any>[]
+    | ExtensionFactory[]
     | (TOptions extends undefined
-        ? () => BlockNoteExtension<any>[]
-        : (options: Partial<TOptions>) => BlockNoteExtension<any>[]),
+        ? () => ExtensionFactory[]
+        : (options: Partial<TOptions>) => ExtensionFactory[]),
 ): (options?: Partial<TOptions>) => BlockSpec<TName, TProps, TContent> {
   return (options = {} as TOptions) => {
     const blockConfig =
diff --git a/packages/core/src/schema/blocks/internal.ts b/packages/core/src/schema/blocks/internal.ts
index 3cfb268d6f..9fda423cd7 100644
--- a/packages/core/src/schema/blocks/internal.ts
+++ b/packages/core/src/schema/blocks/internal.ts
@@ -1,7 +1,7 @@
 import { Attribute, Attributes, Editor, Node } from "@tiptap/core";
 import { defaultBlockToHTML } from "../../blocks/defaultBlockHelpers.js";
 import type { BlockNoteEditor } from "../../editor/BlockNoteEditor.js";
-import { BlockNoteExtension } from "../../editor/BlockNoteExtension.js";
+import type { ExtensionFactory } from "../../editor/BlockNoteExtension.js";
 import { mergeCSSClasses } from "../../util/browser.js";
 import { camelToDataKebab } from "../../util/string.js";
 import { InlineContentSchema } from "../inlineContent/types.js";
@@ -203,7 +203,7 @@ export function createBlockSpecFromTiptapNode<
 >(
   config: T,
   propSchema: P,
-  extensions?: BlockNoteExtension<any>[],
+  extensions?: ExtensionFactory[],
 ): LooseBlockSpec<T["type"], P, T["content"]> {
   return {
     config: {
diff --git a/packages/core/src/schema/blocks/types.ts b/packages/core/src/schema/blocks/types.ts
index 87ec01ab1b..96e2a5984a 100644
--- a/packages/core/src/schema/blocks/types.ts
+++ b/packages/core/src/schema/blocks/types.ts
@@ -4,7 +4,10 @@ import type { Node, NodeViewRendererProps } from "@tiptap/core";
 import type { Fragment, Schema } from "prosemirror-model";
 import type { ViewMutationRecord } from "prosemirror-view";
 import type { BlockNoteEditor } from "../../editor/BlockNoteEditor.js";
-import type { BlockNoteExtension } from "../../editor/BlockNoteExtension.js";
+import type {
+  Extension,
+  ExtensionFactory,
+} from "../../editor/BlockNoteExtension.js";
 import type {
   InlineContent,
   InlineContentSchema,
@@ -98,7 +101,7 @@ export type BlockSpec<
 > = {
   config: BlockConfig<T, PS, C>;
   implementation: BlockImplementation<T, PS, C>;
-  extensions?: BlockNoteExtension<any>[];
+  extensions?: (Extension | ExtensionFactory)[];
 };
 
 /**
@@ -145,7 +148,7 @@ export type LooseBlockSpec<
 
     node: Node;
   };
-  extensions?: BlockNoteExtension<any>[];
+  extensions?: (Extension | ExtensionFactory)[];
 };
 
 // Utility type. For a given object block schema, ensures that the key of each
@@ -197,7 +200,7 @@ export type BlockSpecs = {
           }
         | undefined;
     };
-    extensions?: BlockNoteExtension<any>[];
+    extensions?: BlockSpec<k>["extensions"];
   };
 };
 
@@ -218,7 +221,7 @@ export type BlockSpecsFromSchema<BS extends BlockSchema> = {
       BS[K]["propSchema"],
       BS[K]["content"]
     >;
-    extensions?: BlockNoteExtension<any>[];
+    extensions?: (Extension | ExtensionFactory)[];
   };
 };
 
diff --git a/packages/core/src/schema/schema.ts b/packages/core/src/schema/schema.ts
index 92d0a7ab88..a7a04e93dc 100644
--- a/packages/core/src/schema/schema.ts
+++ b/packages/core/src/schema/schema.ts
@@ -1,5 +1,5 @@
 import { BlockNoteEditor } from "../editor/BlockNoteEditor.js";
-import { createDependencyGraph, toposortReverse } from "../util/topo-sort.js";
+import { sortByDependencies } from "../util/topo-sort.js";
 import {
   BlockNoDefaults,
   BlockSchema,
@@ -80,41 +80,16 @@ export class CustomBlockNoteSchema<
   }
 
   private init() {
-    const dag = createDependencyGraph();
-    const defaultSet = new Set<string>();
-    dag.set("default", defaultSet);
-
-    for (const [key, specDef] of Object.entries({
-      ...this.opts.blockSpecs,
-      ...this.opts.inlineContentSpecs,
-      ...this.opts.styleSpecs,
-    })) {
-      if (specDef.implementation?.runsBefore) {
-        dag.set(key, new Set(specDef.implementation.runsBefore));
-      } else {
-        defaultSet.add(key);
-      }
-    }
-    const sortedSpecs = toposortReverse(dag);
-    const defaultIndex = sortedSpecs.findIndex((set) => set.has("default"));
-
-    /**
-     * The priority of a block is described relative to the "default" block (an arbitrary block which can be used as the reference)
-     *
-     * Since blocks are topologically sorted, we can see what their relative position is to the "default" block
-     * Each layer away from the default block is 10 priority points (arbitrarily chosen)
-     * The default block is fixed at 101 (1 point higher than any tiptap extension, giving priority to custom blocks than any defaults)
-     *
-     * This is a bit of a hack, but it's a simple way to ensure that custom blocks are always rendered with higher priority than default blocks
-     * and that custom blocks are rendered in the order they are defined in the schema
-     */
-    const getPriority = (key: string) => {
-      const index = sortedSpecs.findIndex((set) => set.has(key));
-      // the default index should map to 101
-      // one before the default index is 91
-      // one after is 111
-      return 91 + (index + defaultIndex) * 10;
-    };
+    const getPriority = sortByDependencies(
+      Object.entries({
+        ...this.opts.blockSpecs,
+        ...this.opts.inlineContentSpecs,
+        ...this.opts.styleSpecs,
+      }).map(([key, val]) => ({
+        key: key,
+        runsBefore: val.implementation?.runsBefore ?? [],
+      })),
+    );
 
     const blockSpecs = Object.fromEntries(
       Object.entries(this.opts.blockSpecs).map(([key, blockSpec]) => {
diff --git a/packages/core/src/util/topo-sort.ts b/packages/core/src/util/topo-sort.ts
index fb09fc2075..ea0f7fa891 100644
--- a/packages/core/src/util/topo-sort.ts
+++ b/packages/core/src/util/topo-sort.ts
@@ -158,3 +158,49 @@ export function hasDependency(
   const dependents = graph.get(from);
   return dependents ? dependents.has(to) : false;
 }
+
+/**
+ * Sorts a list of items by their dependencies
+ * @returns A function which can retrieve the priority of an item
+ */
+export function sortByDependencies(
+  items: { key: string; runsBefore?: ReadonlyArray<string> }[],
+) {
+  const dag = createDependencyGraph();
+
+  for (const item of items) {
+    if (Array.isArray(item.runsBefore) && item.runsBefore.length > 0) {
+      item.runsBefore.forEach((runBefore) => {
+        addDependency(dag, item.key, runBefore);
+      });
+    } else {
+      addDependency(dag, "default", item.key);
+    }
+  }
+  const sortedSpecs = toposortReverse(dag);
+  const defaultIndex = sortedSpecs.findIndex((set) => set.has("default"));
+
+  /**
+   * The priority of an item is described relative to the "default" (an arbitrary string which can be used as the reference)
+   *
+   * Since items are topologically sorted, we can see what their relative position is to the "default"
+   * Each layer away from the default is 10 priority points (arbitrarily chosen)
+   * The default is fixed at 101 (1 point higher than any tiptap extension, giving priority to custom blocks than any defaults)
+   *
+   * This is a bit of a hack, but it's a simple way to ensure that custom items are always rendered with higher priority than default items
+   * and that custom items are rendered in the order they are defined in the list
+   */
+
+  /**
+   * Retrieves the priority of an item based on its position in the topologically sorted list
+   * @param key - The key of the item to get the priority of
+   * @returns The priority of the item
+   */
+  return (key: string) => {
+    const index = sortedSpecs.findIndex((set) => set.has(key));
+    // the default index should map to 101
+    // one before the default index is 91
+    // one after is 111
+    return 91 + (index + defaultIndex) * 10;
+  };
+}
diff --git a/packages/core/vite.config.ts b/packages/core/vite.config.ts
index de4bd587d9..505c58a2b4 100644
--- a/packages/core/vite.config.ts
+++ b/packages/core/vite.config.ts
@@ -1,6 +1,6 @@
 import * as path from "path";
 import { webpackStats } from "rollup-plugin-webpack-stats";
-import { defineConfig } from "vite";
+import { defineConfig } from "vitest/config";
 import pkg from "./package.json";
 // import eslintPlugin from "vite-plugin-eslint";
 
@@ -21,6 +21,7 @@ export default defineConfig({
         comments: path.resolve(__dirname, "src/comments/index.ts"),
         blocks: path.resolve(__dirname, "src/blocks/index.ts"),
         locales: path.resolve(__dirname, "src/i18n/index.ts"),
+        extensions: path.resolve(__dirname, "src/extensions/index.ts"),
         yjs: path.resolve(__dirname, "src/yjs/index.ts"),
       },
       name: "blocknote",
diff --git a/packages/react/package.json b/packages/react/package.json
index 142ef03d43..0f5b8c5587 100644
--- a/packages/react/package.json
+++ b/packages/react/package.json
@@ -61,12 +61,17 @@
     "@blocknote/core": "0.42.3",
     "@emoji-mart/data": "^1.2.1",
     "@floating-ui/react": "^0.27.16",
-    "@tiptap/core": "^3.11.0",
-    "@tiptap/pm": "^3.11.0",
-    "@tiptap/react": "^3.11.0",
+    "@floating-ui/utils": "0.2.10",
+    "@tanstack/react-store": "0.7.7",
+    "@tiptap/core": "^3.7.2",
+    "@tiptap/pm": "^3.7.2",
+    "@tiptap/react": "^3.7.2",
+    "@types/use-sync-external-store": "1.5.0",
     "emoji-mart": "^5.6.0",
+    "fast-deep-equal": "^3.1.3",
     "lodash.merge": "^4.6.2",
-    "react-icons": "^5.5.0"
+    "react-icons": "^5.5.0",
+    "use-sync-external-store": "1.6.0"
   },
   "devDependencies": {
     "@types/emoji-mart": "^3.0.14",
diff --git a/packages/react/src/blocks/File/helpers/render/AddFileButton.tsx b/packages/react/src/blocks/File/helpers/render/AddFileButton.tsx
index 44c0fc9a23..50e6e21463 100644
--- a/packages/react/src/blocks/File/helpers/render/AddFileButton.tsx
+++ b/packages/react/src/blocks/File/helpers/render/AddFileButton.tsx
@@ -1,9 +1,11 @@
 import { FileBlockConfig } from "@blocknote/core";
+import { FilePanel } from "@blocknote/core/extensions";
 import { ReactNode, useCallback } from "react";
 import { RiFile2Line } from "react-icons/ri";
 
 import { useDictionary } from "../../../../i18n/dictionary.js";
 import { ReactCustomBlockRenderProps } from "../../../../schema/ReactBlockSpec.js";
+import { useExtension } from "../../../../hooks/useExtension.js";
 
 export const AddFileButton = (
   props: Omit<
@@ -19,6 +21,8 @@ export const AddFileButton = (
 ) => {
   const dict = useDictionary();
 
+  const filePanel = useExtension(FilePanel);
+
   // Prevents focus from moving to the button.
   const addFileButtonMouseDownHandler = useCallback(
     (event: React.MouseEvent) => {
@@ -28,12 +32,8 @@ export const AddFileButton = (
   );
   // Opens the file toolbar.
   const addFileButtonClickHandler = useCallback(() => {
-    props.editor.transact((tr) =>
-      tr.setMeta(props.editor.filePanel!.plugins[0], {
-        block: props.block,
-      }),
-    );
-  }, [props.block, props.editor]);
+    props.editor.transact(() => filePanel.showMenu(props.block.id));
+  }, [filePanel, props.block.id, props.editor]);
 
   return (
     <div
diff --git a/packages/react/src/blocks/ToggleWrapper/ToggleWrapper.tsx b/packages/react/src/blocks/ToggleWrapper/ToggleWrapper.tsx
index 82f29d638f..f688b37922 100644
--- a/packages/react/src/blocks/ToggleWrapper/ToggleWrapper.tsx
+++ b/packages/react/src/blocks/ToggleWrapper/ToggleWrapper.tsx
@@ -1,12 +1,13 @@
 import {
   Block,
+  blockHasType,
   defaultToggledState,
   UnreachableCaseError,
 } from "@blocknote/core";
-import { ReactNode, useReducer, useState } from "react";
+import { ReactNode, useReducer } from "react";
 
 import { ReactCustomBlockRenderProps } from "../../schema/ReactBlockSpec.js";
-import { useEditorChange } from "../../hooks/useEditorChange.js";
+import { useEditorState } from "../../hooks/useEditorState.js";
 
 const showChildrenReducer = (
   showChildren: boolean,
@@ -51,7 +52,6 @@ export const ToggleWrapper = (
     showChildrenReducer,
     (toggledState || defaultToggledState).get(block),
   );
-  const [childCount, setChildCount] = useState(block.children.length);
 
   const handleToggle = (block: Block<any, any, any>) => {
     (toggledState || defaultToggledState).set(
@@ -77,28 +77,34 @@ export const ToggleWrapper = (
     });
   };
 
-  useEditorChange(() => {
-    if ("isToggleable" in block.props && !block.props.isToggleable) {
-      return;
-    }
-
-    const newBlock = editor.getBlock(block)!;
-    const newChildCount = newBlock.children.length ?? 0;
-
-    if (newChildCount > childCount) {
-      // If a child block is added while children are hidden, show children.
-      if (!showChildren) {
-        handleChildAdded(newBlock);
+  const childCount = useEditorState({
+    editor,
+    selector: ({ editor }) => {
+      if (
+        !blockHasType(block, editor, block.type, { isToggleable: "boolean" }) &&
+        !block.props.isToggleable
+      ) {
+        return 0;
       }
-    } else if (newChildCount === 0 && newChildCount < childCount) {
-      // If the last child block is removed while children are shown, hide
-      // children.
-      if (showChildren) {
-        handleLastChildRemoved(newBlock);
+
+      const newBlock = editor.getBlock(block)!;
+      const newChildCount = newBlock.children.length || 0;
+
+      if (newChildCount > childCount) {
+        // If a child block is added while children are hidden, show children.
+        if (!showChildren) {
+          handleChildAdded(newBlock);
+        }
+      } else if (newChildCount === 0 && newChildCount < childCount) {
+        // If the last child block is removed while children are shown, hide
+        // children.
+        if (showChildren) {
+          handleLastChildRemoved(newBlock);
+        }
       }
-    }
 
-    setChildCount(newChildCount);
+      return newChildCount;
+    },
   });
 
   if ("isToggleable" in block.props && !block.props.isToggleable) {
diff --git a/packages/react/src/components/Comments/Comment.tsx b/packages/react/src/components/Comments/Comment.tsx
index 55c3e1af23..2aef7f6d90 100644
--- a/packages/react/src/components/Comments/Comment.tsx
+++ b/packages/react/src/components/Comments/Comment.tsx
@@ -1,6 +1,7 @@
 "use client";
 
 import { mergeCSSClasses } from "@blocknote/core";
+import { Comments } from "@blocknote/core/extensions";
 import type { CommentData, ThreadData } from "@blocknote/core/comments";
 import { MouseEvent, ReactNode, useCallback, useState } from "react";
 import {
@@ -13,8 +14,8 @@ import {
 } from "react-icons/ri";
 
 import { useComponentsContext } from "../../editor/ComponentsContext.js";
-import { useBlockNoteEditor } from "../../hooks/useBlockNoteEditor.js";
 import { useCreateBlockNote } from "../../hooks/useCreateBlockNote.js";
+import { useExtension } from "../../hooks/useExtension.js";
 import { useDictionary } from "../../i18n/dictionary.js";
 import { CommentEditor } from "./CommentEditor.js";
 import { EmojiPicker } from "./EmojiPicker.js";
@@ -42,12 +43,7 @@ export const Comment = ({
 }: CommentProps) => {
   // TODO: if REST API becomes popular, all interactions (click handlers) should implement a loading state and error state
   // (or optimistic local updates)
-
-  const editor = useBlockNoteEditor();
-
-  if (!editor.comments) {
-    throw new Error("Comments plugin not found");
-  }
+  const comments = useExtension(Comments);
 
   const dict = useDictionary();
 
@@ -60,7 +56,7 @@ export const Comment = ({
         emptyDocument: dict.placeholders.edit_comment,
       },
     },
-    schema: editor.comments.commentEditorSchema || defaultCommentEditorSchema,
+    schema: comments.commentEditorSchema || defaultCommentEditorSchema,
   });
 
   const Components = useComponentsContext()!;
@@ -68,11 +64,7 @@ export const Comment = ({
   const [isEditing, setEditing] = useState(false);
   const [emojiPickerOpen, setEmojiPickerOpen] = useState(false);
 
-  if (!editor.comments) {
-    throw new Error("Comments plugin not found");
-  }
-
-  const threadStore = editor.comments.threadStore;
+  const threadStore = comments.threadStore;
 
   const handleEdit = useCallback(() => {
     setEditing(true);
@@ -136,7 +128,7 @@ export const Comment = ({
     });
   }, [thread.id, threadStore]);
 
-  const user = useUser(editor, comment.userId);
+  const user = useUser(comment.userId);
 
   if (!comment.body) {
     return null;
diff --git a/packages/react/src/components/Comments/CommentEditor.tsx b/packages/react/src/components/Comments/CommentEditor.tsx
index 4b22eb6c5c..6030cb528e 100644
--- a/packages/react/src/components/Comments/CommentEditor.tsx
+++ b/packages/react/src/components/Comments/CommentEditor.tsx
@@ -1,7 +1,7 @@
 import { BlockNoteEditor } from "@blocknote/core";
 import { FC, useCallback, useEffect, useState } from "react";
 import { useComponentsContext } from "../../editor/ComponentsContext.js";
-import { useEditorChange } from "../../hooks/useEditorChange.js";
+import { useEditorState } from "../../hooks/useEditorState.js";
 
 /**
  * The CommentEditor component displays an editor for creating or editing a comment.
@@ -23,14 +23,13 @@ export const CommentEditor = (props: {
   editor: BlockNoteEditor<any, any, any>;
 }) => {
   const [isFocused, setIsFocused] = useState(false);
-  const [isEmpty, setIsEmpty] = useState(props.editor.isEmpty);
+  const isEmpty = useEditorState({
+    editor: props.editor,
+    selector: ({ editor }) => editor.isEmpty,
+  });
 
   const components = useComponentsContext()!;
 
-  useEditorChange(() => {
-    setIsEmpty(props.editor.isEmpty);
-  }, props.editor);
-
   const onFocus = useCallback(() => {
     setIsFocused(true);
   }, []);
diff --git a/packages/react/src/components/Comments/Comments.tsx b/packages/react/src/components/Comments/Comments.tsx
index ab6203a6f6..7e375094cb 100644
--- a/packages/react/src/components/Comments/Comments.tsx
+++ b/packages/react/src/components/Comments/Comments.tsx
@@ -3,7 +3,6 @@ import { ThreadData } from "@blocknote/core/comments";
 import { useComponentsContext } from "../../editor/ComponentsContext.js";
 import { useDictionary } from "../../i18n/dictionary.js";
 import { Comment } from "./Comment.js";
-import { useBlockNoteEditor } from "../../hooks/useBlockNoteEditor.js";
 import { useUsers } from "./useUsers.js";
 
 export type CommentsProps = {
@@ -17,8 +16,7 @@ export const Comments = ({
 }: CommentsProps) => {
   const Components = useComponentsContext()!;
   const dict = useDictionary();
-  const editor = useBlockNoteEditor();
-  const users = useUsers(editor, thread.resolvedBy ? [thread.resolvedBy] : []);
+  const users = useUsers(thread.resolvedBy ? [thread.resolvedBy] : []);
 
   // Maps all comments to elements.
   const comments = thread.comments.map((comment, index) => (
diff --git a/packages/react/src/components/Comments/FloatingComposer.tsx b/packages/react/src/components/Comments/FloatingComposer.tsx
index e756f902b0..546bb8caa2 100644
--- a/packages/react/src/components/Comments/FloatingComposer.tsx
+++ b/packages/react/src/components/Comments/FloatingComposer.tsx
@@ -1,8 +1,9 @@
 import { mergeCSSClasses } from "@blocknote/core";
+import { Comments } from "@blocknote/core/extensions";
 
 import { useComponentsContext } from "../../editor/ComponentsContext.js";
-import { useBlockNoteEditor } from "../../hooks/useBlockNoteEditor.js";
 import { useCreateBlockNote } from "../../hooks/useCreateBlockNote.js";
+import { useExtension } from "../../hooks/useExtension.js";
 import { useDictionary } from "../../i18n/dictionary.js";
 import { CommentEditor } from "./CommentEditor.js";
 import { defaultCommentEditorSchema } from "./defaultCommentEditorSchema.js";
@@ -13,13 +14,7 @@ import { defaultCommentEditorSchema } from "./defaultCommentEditorSchema.js";
  * It's used when the user highlights a parts of the document to create a new comment / thread.
  */
 export function FloatingComposer() {
-  const editor = useBlockNoteEditor();
-
-  if (!editor.comments) {
-    throw new Error("Comments plugin not found");
-  }
-
-  const comments = editor.comments;
+  const comments = useExtension(Comments);
 
   const Components = useComponentsContext()!;
   const dict = useDictionary();
@@ -32,7 +27,7 @@ export function FloatingComposer() {
         emptyDocument: dict.placeholders.new_comment,
       },
     },
-    schema: editor.comments.commentEditorSchema || defaultCommentEditorSchema,
+    schema: comments.commentEditorSchema || defaultCommentEditorSchema,
   });
 
   return (
diff --git a/packages/react/src/components/Comments/FloatingComposerController.tsx b/packages/react/src/components/Comments/FloatingComposerController.tsx
index aa50f5b684..f26dc8eae0 100644
--- a/packages/react/src/components/Comments/FloatingComposerController.tsx
+++ b/packages/react/src/components/Comments/FloatingComposerController.tsx
@@ -6,14 +6,16 @@ import {
   InlineContentSchema,
   StyleSchema,
 } from "@blocknote/core";
-import { UseFloatingOptions, flip, offset, shift } from "@floating-ui/react";
-import { ComponentProps, FC, useEffect } from "react";
+import { Comments } from "@blocknote/core/extensions";
+import { flip, offset, shift } from "@floating-ui/react";
+import { ComponentProps, FC, useEffect, useMemo, useState } from "react";
 
 import { useBlockNoteEditor } from "../../hooks/useBlockNoteEditor.js";
-import { useEditorSelectionBoundingBox } from "../../hooks/useEditorSelectionBoundingBox.js";
-import { useUIElementPositioning } from "../../hooks/useUIElementPositioning.js";
-import { useUIPluginState } from "../../hooks/useUIPluginState.js";
+import { useExtension, useExtensionState } from "../../hooks/useExtension.js";
 import { FloatingComposer } from "./FloatingComposer.js";
+import { PositionPopover } from "../Popovers/PositionPopover.js";
+import { useEditorState } from "../../hooks/useEditorState.js";
+import { FloatingUIOptions } from "../Popovers/FloatingUIOptions.js";
 
 export const FloatingComposerController = <
   B extends BlockSchema = DefaultBlockSchema,
@@ -21,52 +23,69 @@ export const FloatingComposerController = <
   S extends StyleSchema = DefaultStyleSchema,
 >(props: {
   floatingComposer?: FC<ComponentProps<typeof FloatingComposer>>;
-  floatingOptions?: Partial<UseFloatingOptions>;
+  floatingUIOptions?: FloatingUIOptions;
 }) => {
+  const [open, setOpen] = useState(false);
+
   const editor = useBlockNoteEditor<B, I, S>();
 
-  if (!editor.comments) {
-    throw new Error(
-      "FloatingComposerController can only be used when BlockNote editor has enabled comments",
-    );
-  }
+  const comments = useExtension(Comments);
 
-  const comments = editor.comments;
-  useEffect(() => {
-    const offUpdate = comments.onUpdate((state) =>
-      editor.setForceSelectionVisible(state.pendingComment),
-    );
+  // TODO: `setForceSelectionVisible` no longer exists?
+  // useEffect(() => {
+  //   const offUpdate = comments.onUpdate((state) =>
+  //     editor.setForceSelectionVisible(state.pendingComment),
+  //   );
+
+  //   return () => offUpdate();
+  // }, [comments, editor]);
 
-    return () => offUpdate();
-  }, [comments, editor]);
+  const pendingComment = useExtensionState(Comments, {
+    editor,
+    selector: (state) => state.pendingComment,
+  });
+  useEffect(() => {
+    setOpen(pendingComment);
+  }, [pendingComment]);
 
-  const state = useUIPluginState(comments.onUpdate.bind(comments));
+  const position = useEditorState({
+    editor,
+    selector: ({ editor }) =>
+      pendingComment
+        ? {
+            from: editor.prosemirrorState.selection.from,
+            to: editor.prosemirrorState.selection.to,
+          }
+        : undefined,
+  });
 
-  // the reference position is updated automatically when the selection moves,
-  // this can happen when the document is updated by a remote user
-  const referencePos = useEditorSelectionBoundingBox(state?.pendingComment);
+  const floatingUIOptions = useMemo<FloatingUIOptions>(
+    () => ({
+      useFloatingOptions: {
+        open,
+        // Needed as hooks like `useDismiss` call `onOpenChange` to change the
+        // open state.
+        onOpenChange: (open) => {
+          if (!open) {
+            comments.stopPendingComment();
+            editor.focus();
+          }
 
-  const { isMounted, ref, style, getFloatingProps } = useUIElementPositioning(
-    state?.pendingComment || false,
-    referencePos || null,
-    5000,
-    {
-      placement: "bottom",
-      middleware: [offset(10), shift(), flip()],
-      onOpenChange: (open) => {
-        if (!open) {
-          comments.stopPendingComment();
-          editor.focus();
-        }
+          setOpen(open);
+        },
+        placement: "bottom",
+        middleware: [offset(10), shift(), flip()],
+      },
+      elementProps: {
+        style: {
+          zIndex: 60,
+        },
       },
-      ...props.floatingOptions,
-    },
+      ...props.floatingUIOptions,
+    }),
+    [comments, editor, open, props.floatingUIOptions],
   );
 
-  if (!isMounted || !state) {
-    return null;
-  }
-
   // nice to have improvements would be:
   // - transition transform property so composer box animates when remote document is changed
   // - fade out on close
@@ -74,8 +93,8 @@ export const FloatingComposerController = <
   const Component = props.floatingComposer || FloatingComposer;
 
   return (
-    <div ref={ref} style={style} {...getFloatingProps()}>
+    <PositionPopover position={position} {...floatingUIOptions}>
       <Component />
-    </div>
+    </PositionPopover>
   );
 };
diff --git a/packages/react/src/components/Comments/FloatingThreadController.tsx b/packages/react/src/components/Comments/FloatingThreadController.tsx
index 3df05ef6b0..c0070ffc46 100644
--- a/packages/react/src/components/Comments/FloatingThreadController.tsx
+++ b/packages/react/src/components/Comments/FloatingThreadController.tsx
@@ -1,29 +1,11 @@
-import {
-  BlockSchema,
-  DefaultBlockSchema,
-  DefaultInlineContentSchema,
-  DefaultStyleSchema,
-  InlineContentSchema,
-  StyleSchema,
-} from "@blocknote/core";
-import {
-  UseFloatingOptions,
-  autoUpdate,
-  flip,
-  offset,
-  shift,
-} from "@floating-ui/react";
-import {
-  ComponentProps,
-  FC,
-  useCallback,
-  useEffect,
-  useLayoutEffect,
-} from "react";
+import { Comments } from "@blocknote/core/extensions";
+import { flip, offset, shift } from "@floating-ui/react";
+import { ComponentProps, FC, useEffect, useMemo, useState } from "react";
 
 import { useBlockNoteEditor } from "../../hooks/useBlockNoteEditor.js";
-import { useUIElementPositioning } from "../../hooks/useUIElementPositioning.js";
-import { useUIPluginState } from "../../hooks/useUIPluginState.js";
+import { useExtensionState } from "../../hooks/useExtension.js";
+import { FloatingUIOptions } from "../Popovers/FloatingUIOptions.js";
+import { PositionPopover } from "../Popovers/PositionPopover.js";
 import { Thread } from "./Thread.js";
 import { useThreads } from "./useThreads.js";
 
@@ -31,78 +13,49 @@ import { useThreads } from "./useThreads.js";
  * This component is used to display a thread in a floating card.
  * It can be used when the user clicks on a thread / comment in the document.
  */
-export const FloatingThreadController = <
-  B extends BlockSchema = DefaultBlockSchema,
-  I extends InlineContentSchema = DefaultInlineContentSchema,
-  S extends StyleSchema = DefaultStyleSchema,
->(props: {
+export const FloatingThreadController = (props: {
   floatingThread?: FC<ComponentProps<typeof Thread>>;
-  floatingOptions?: Partial<UseFloatingOptions>;
+  floatingUIOptions?: FloatingUIOptions;
 }) => {
-  const editor = useBlockNoteEditor<B, I, S>();
-
-  if (!editor.comments) {
-    throw new Error(
-      "FloatingComposerController can only be used when BlockNote editor has enabled comments",
-    );
-  }
-
-  const state = useUIPluginState(
-    editor.comments.onUpdate.bind(editor.comments),
-  );
-
-  const { isMounted, ref, style, getFloatingProps, setReference } =
-    useUIElementPositioning(!!state?.selectedThreadId, null, 5000, {
-      placement: "bottom",
-      middleware: [offset(10), shift(), flip()],
-      onOpenChange: (open) => {
-        if (!open) {
-          editor.comments?.selectThread(undefined);
-          editor.focus();
-        }
-      },
-      whileElementsMounted: autoUpdate,
-      ...props.floatingOptions,
-    });
-
-  // Positioning with [data-bn-thread-id] attribute is a bit hacky,
-  // we could probably also use the thread position from the plugin state?
-  // for now, this works ok
-  const updateRef = useCallback(() => {
-    if (!state?.selectedThreadId) {
-      return;
-    }
-
-    const el = editor.domElement?.querySelector(
-      `[data-bn-thread-id="${state?.selectedThreadId}"]`,
-    );
-    if (el) {
-      setReference(el);
-    }
-  }, [setReference, editor, state?.selectedThreadId]);
-
-  // Remote cursor updates and other edits can cause the ref to break
+  const editor = useBlockNoteEditor<any, any, any>();
+
+  const [open, setOpen] = useState(false);
+
+  const selectedThread = useExtensionState(Comments, {
+    editor,
+    selector: (state) =>
+      state.selectedThreadId
+        ? {
+            id: state.selectedThreadId,
+            position: state.threadPositions.get(state.selectedThreadId),
+          }
+        : undefined,
+  });
   useEffect(() => {
-    if (!state?.selectedThreadId) {
-      return;
-    }
-
-    return editor.onChange(() => {
-      updateRef();
-    });
-  }, [editor, updateRef, state?.selectedThreadId]);
-
-  useLayoutEffect(updateRef, [updateRef]);
-
-  const threads = useThreads(editor);
-
-  if (!isMounted || !state) {
-    return null;
-  }
-
-  if (!state.selectedThreadId) {
-    return null;
-  }
+    setOpen(!!selectedThread);
+  }, [selectedThread]);
+
+  const threads = useThreads();
+
+  const floatingUIOptions = useMemo<FloatingUIOptions>(
+    () => ({
+      useFloatingOptions: {
+        open,
+        // Needed as hooks like `useDismiss` call `onOpenChange` to change the
+        // open state.
+        onOpenChange: setOpen,
+        placement: "bottom",
+        middleware: [offset(10), shift(), flip()],
+      },
+      elementProps: {
+        style: {
+          zIndex: 30,
+        },
+      },
+      ...props.floatingUIOptions,
+    }),
+    [open, props.floatingUIOptions],
+  );
 
   // nice to have improvements:
   // - fade out on close
@@ -111,11 +64,10 @@ export const FloatingThreadController = <
   const Component = props.floatingThread || Thread;
 
   return (
-    <div ref={ref} style={style} {...getFloatingProps()}>
-      <Component
-        thread={threads.get(state.selectedThreadId)!}
-        selected={true}
-      />
-    </div>
+    <PositionPopover position={selectedThread?.position} {...floatingUIOptions}>
+      {selectedThread && (
+        <Component thread={threads.get(selectedThread.id)!} selected={true} />
+      )}
+    </PositionPopover>
   );
 };
diff --git a/packages/react/src/components/Comments/ReactionBadge.tsx b/packages/react/src/components/Comments/ReactionBadge.tsx
index 0a70b4283e..15a7c1cb15 100644
--- a/packages/react/src/components/Comments/ReactionBadge.tsx
+++ b/packages/react/src/components/Comments/ReactionBadge.tsx
@@ -1,11 +1,12 @@
 import { mergeCSSClasses } from "@blocknote/core";
+import { Comments } from "@blocknote/core/extensions";
 import { CommentData } from "@blocknote/core/comments";
 import { useState } from "react";
 
 import { useDictionary } from "../../i18n/dictionary.js";
 import { useComponentsContext } from "../../editor/ComponentsContext.js";
-import { useBlockNoteEditor } from "../../hooks/useBlockNoteEditor.js";
 import { useUsers } from "./useUsers.js";
+import { useExtension } from "../../hooks/useExtension.js";
 
 export const ReactionBadge = (props: {
   comment: CommentData;
@@ -15,12 +16,7 @@ export const ReactionBadge = (props: {
   const Components = useComponentsContext()!;
   const dict = useDictionary();
 
-  const editor = useBlockNoteEditor();
-  if (!editor.comments) {
-    throw new Error(
-      "ReactionBadge must be used inside a BlockNote editor with comments enabled",
-    );
-  }
+  const comments = useExtension(Comments);
 
   const reaction = props.comment.reactions.find(
     (reaction) => reaction.emoji === props.emoji,
@@ -32,7 +28,7 @@ export const ReactionBadge = (props: {
   }
 
   const [userIds, setUserIds] = useState<string[]>([]);
-  const users = useUsers(editor, userIds);
+  const users = useUsers(userIds);
 
   return (
     <Components.Generic.Badge.Root
@@ -40,7 +36,7 @@ export const ReactionBadge = (props: {
       className={mergeCSSClasses("bn-badge", "bn-comment-reaction")}
       text={reaction.userIds.length.toString()}
       icon={reaction.emoji}
-      isSelected={editor.comments.threadStore.auth.canDeleteReaction(
+      isSelected={comments.threadStore.auth.canDeleteReaction(
         props.comment,
         reaction.emoji,
       )}
diff --git a/packages/react/src/components/Comments/Thread.tsx b/packages/react/src/components/Comments/Thread.tsx
index 3ad39cd306..40d4cbd096 100644
--- a/packages/react/src/components/Comments/Thread.tsx
+++ b/packages/react/src/components/Comments/Thread.tsx
@@ -1,14 +1,15 @@
 import { mergeCSSClasses } from "@blocknote/core";
+import { Comments as CommentsExtension } from "@blocknote/core/extensions";
 import { ThreadData } from "@blocknote/core/comments";
 import { FocusEvent, useCallback } from "react";
 
 import { useComponentsContext } from "../../editor/ComponentsContext.js";
-import { useBlockNoteEditor } from "../../hooks/useBlockNoteEditor.js";
 import { useCreateBlockNote } from "../../hooks/useCreateBlockNote.js";
 import { useDictionary } from "../../i18n/dictionary.js";
 import { CommentEditor } from "./CommentEditor.js";
 import { Comments } from "./Comments.js";
 import { defaultCommentEditorSchema } from "./defaultCommentEditorSchema.js";
+import { useExtension } from "../../hooks/useExtension.js";
 
 export type ThreadProps = {
   /**
@@ -66,12 +67,7 @@ export const Thread = ({
   const Components = useComponentsContext()!;
   const dict = useDictionary();
 
-  const editor = useBlockNoteEditor();
-
-  const comments = editor.comments;
-  if (!comments) {
-    throw new Error("Comments plugin not found");
-  }
+  const comments = useExtension(CommentsExtension);
 
   const newCommentEditor = useCreateBlockNote({
     trailingBlock: false,
@@ -81,7 +77,7 @@ export const Thread = ({
         emptyDocument: dict.placeholders.comment_reply,
       },
     },
-    schema: editor.comments.commentEditorSchema || defaultCommentEditorSchema,
+    schema: comments.commentEditorSchema || defaultCommentEditorSchema,
   });
 
   const onNewCommentSave = useCallback(async () => {
diff --git a/packages/react/src/components/Comments/ThreadsSidebar.tsx b/packages/react/src/components/Comments/ThreadsSidebar.tsx
index 5afd7a7e04..e56d084ee4 100644
--- a/packages/react/src/components/Comments/ThreadsSidebar.tsx
+++ b/packages/react/src/components/Comments/ThreadsSidebar.tsx
@@ -1,8 +1,9 @@
 import { BlockNoteEditor, UnreachableCaseError } from "@blocknote/core";
+import { Comments } from "@blocknote/core/extensions";
 import { ThreadData } from "@blocknote/core/comments";
 import React, { FocusEvent, useCallback, useMemo } from "react";
 import { useBlockNoteEditor } from "../../hooks/useBlockNoteEditor.js";
-import { useUIPluginState } from "../../hooks/useUIPluginState.js";
+import { useExtension, useExtensionState } from "../../hooks/useExtension.js";
 import { Thread } from "./Thread.js";
 import { useThreads } from "./useThreads.js";
 
@@ -22,10 +23,11 @@ const ThreadItem = React.memo(
   ({
     thread,
     selectedThreadId,
-    editor,
     maxCommentsBeforeCollapse,
     referenceText,
   }: ThreadItemProps) => {
+    const comments = useExtension(Comments);
+
     const onFocus = useCallback(
       (event: FocusEvent) => {
         // If the focused element is within the action toolbar, we don't want to
@@ -34,9 +36,9 @@ const ThreadItem = React.memo(
           return;
         }
 
-        editor.comments?.selectThread(thread.id);
+        comments.selectThread(thread.id);
       },
-      [editor.comments, thread.id],
+      [comments, thread.id],
     );
 
     const onBlur = useCallback(
@@ -64,10 +66,10 @@ const ThreadItem = React.memo(
           !parentThreadElement ||
           !parentThreadElement.contains(targetElement)
         ) {
-          editor.comments?.selectThread(undefined);
+          comments.selectThread(undefined);
         }
       },
-      [editor.comments],
+      [comments],
     );
 
     return (
@@ -188,22 +190,11 @@ export function ThreadsSidebar(props: {
    */
   sort?: "position" | "recent-activity" | "oldest";
 }) {
-  const editor = useBlockNoteEditor();
-
-  if (!editor.comments) {
-    throw new Error("Comments plugin not found");
-  }
-
-  // Note: because "threadPositions" is part of the state,
-  // this will potentially trigger a re-render on every document update
-  // this means we need to be mindful of children memoization
-  const state = useUIPluginState(
-    editor.comments.onUpdate.bind(editor.comments),
-  );
+  const editor = useBlockNoteEditor<any, any, any>();
 
-  const selectedThreadId = state?.selectedThreadId;
+  const { selectedThreadId, threadPositions } = useExtensionState(Comments);
 
-  const threads = useThreads(editor);
+  const threads = useThreads();
 
   const filteredAndSortedThreads = useMemo(() => {
     const threadsArray = Array.from(threads.values());
@@ -211,7 +202,7 @@ export function ThreadsSidebar(props: {
     const sortedThreads = sortThreads(
       threadsArray,
       props.sort || "position",
-      state?.threadPositions,
+      threadPositions,
     );
 
     const ret: Array<{ thread: ThreadData; referenceText: string }> = [];
@@ -223,7 +214,7 @@ export function ThreadsSidebar(props: {
             thread,
             referenceText: getReferenceText(
               editor,
-              state?.threadPositions.get(thread.id),
+              threadPositions.get(thread.id),
             ),
           });
         }
@@ -233,7 +224,7 @@ export function ThreadsSidebar(props: {
             thread,
             referenceText: getReferenceText(
               editor,
-              state?.threadPositions.get(thread.id),
+              threadPositions.get(thread.id),
             ),
           });
         }
@@ -241,7 +232,7 @@ export function ThreadsSidebar(props: {
     }
 
     return ret;
-  }, [threads, state?.threadPositions, props.filter, props.sort, editor]);
+  }, [threads, props.sort, props.filter, threadPositions, editor]);
 
   return (
     <div className={"bn-threads-sidebar"}>
diff --git a/packages/react/src/components/Comments/useThreads.ts b/packages/react/src/components/Comments/useThreads.ts
index 3977827c71..cc65d634a5 100644
--- a/packages/react/src/components/Comments/useThreads.ts
+++ b/packages/react/src/components/Comments/useThreads.ts
@@ -1,15 +1,14 @@
-import { BlockNoteEditor } from "@blocknote/core";
+import { Comments } from "@blocknote/core/extensions";
 import { ThreadData } from "@blocknote/core/comments";
 import { useCallback, useRef, useSyncExternalStore } from "react";
 
+import { useExtension } from "../../hooks/useExtension.js";
+
 /**
  * Bridges the ThreadStore to React using useSyncExternalStore.
  */
-export function useThreads(editor: BlockNoteEditor<any, any, any>) {
-  const comments = editor.comments;
-  if (!comments) {
-    throw new Error("Comments plugin not found");
-  }
+export function useThreads() {
+  const comments = useExtension(Comments);
 
   const store = comments.threadStore;
 
diff --git a/packages/react/src/components/Comments/useUsers.ts b/packages/react/src/components/Comments/useUsers.ts
index 13d9568e03..27c8c0b8e8 100644
--- a/packages/react/src/components/Comments/useUsers.ts
+++ b/packages/react/src/components/Comments/useUsers.ts
@@ -1,25 +1,18 @@
-import { BlockNoteEditor } from "@blocknote/core";
+import { Comments } from "@blocknote/core/extensions";
 import { User } from "@blocknote/core/comments";
 import { useCallback, useMemo, useSyncExternalStore } from "react";
 
-export function useUser(
-  editor: BlockNoteEditor<any, any, any>,
-  userId: string,
-) {
-  return useUsers(editor, [userId]).get(userId);
+import { useExtension } from "../../hooks/useExtension.js";
+
+export function useUser(userId: string) {
+  return useUsers([userId]).get(userId);
 }
 
 /**
  * Bridges the UserStore to React using useSyncExternalStore.
  */
-export function useUsers(
-  editor: BlockNoteEditor<any, any, any>,
-  userIds: string[],
-) {
-  const comments = editor.comments;
-  if (!comments) {
-    throw new Error("Comments plugin not found");
-  }
+export function useUsers(userIds: string[]) {
+  const comments = useExtension(Comments);
 
   const store = comments.userStore;
 
diff --git a/packages/react/src/components/FilePanel/DefaultTabs/EmbedTab.tsx b/packages/react/src/components/FilePanel/DefaultTabs/EmbedTab.tsx
index db07823b75..3d8090bc20 100644
--- a/packages/react/src/components/FilePanel/DefaultTabs/EmbedTab.tsx
+++ b/packages/react/src/components/FilePanel/DefaultTabs/EmbedTab.tsx
@@ -19,15 +19,15 @@ export const EmbedTab = <
   I extends InlineContentSchema = DefaultInlineContentSchema,
   S extends StyleSchema = DefaultStyleSchema,
 >(
-  props: FilePanelProps<I, S>,
+  props: FilePanelProps,
 ) => {
   const Components = useComponentsContext()!;
   const dict = useDictionary();
 
-  const { block } = props;
-
   const editor = useBlockNoteEditor<B, I, S>();
 
+  const block = editor.getBlock(props.blockId)!;
+
   const [currentURL, setCurrentURL] = useState<string>("");
 
   const handleURLChange = useCallback(
@@ -41,7 +41,7 @@ export const EmbedTab = <
     (event: KeyboardEvent) => {
       if (event.key === "Enter") {
         event.preventDefault();
-        editor.updateBlock(block, {
+        editor.updateBlock(block.id, {
           props: {
             name: filenameFromURL(currentURL),
             url: currentURL,
@@ -49,17 +49,17 @@ export const EmbedTab = <
         });
       }
     },
-    [editor, block, currentURL],
+    [editor, block.id, currentURL],
   );
 
   const handleURLClick = useCallback(() => {
-    editor.updateBlock(block, {
+    editor.updateBlock(block.id, {
       props: {
         name: filenameFromURL(currentURL),
         url: currentURL,
       } as any,
     });
-  }, [editor, block, currentURL]);
+  }, [editor, block.id, currentURL]);
 
   return (
     <Components.FilePanel.TabPanel className={"bn-tab-panel"}>
diff --git a/packages/react/src/components/FilePanel/DefaultTabs/UploadTab.tsx b/packages/react/src/components/FilePanel/DefaultTabs/UploadTab.tsx
index 23847eb572..64d5a1c74f 100644
--- a/packages/react/src/components/FilePanel/DefaultTabs/UploadTab.tsx
+++ b/packages/react/src/components/FilePanel/DefaultTabs/UploadTab.tsx
@@ -18,17 +18,19 @@ export const UploadTab = <
   I extends InlineContentSchema = DefaultInlineContentSchema,
   S extends StyleSchema = DefaultStyleSchema,
 >(
-  props: FilePanelProps<I, S> & {
+  props: FilePanelProps & {
     setLoading: (loading: boolean) => void;
   },
 ) => {
   const Components = useComponentsContext()!;
   const dict = useDictionary();
 
-  const { block, setLoading } = props;
+  const { setLoading } = props;
 
   const editor = useBlockNoteEditor<B, I, S>();
 
+  const block = editor.getBlock(props.blockId)!;
+
   const [uploadFailed, setUploadFailed] = useState<boolean>(false);
 
   useEffect(() => {
@@ -50,7 +52,7 @@ export const UploadTab = <
 
         if (editor.uploadFile !== undefined) {
           try {
-            let updateData = await editor.uploadFile(file, block.id);
+            let updateData = await editor.uploadFile(file, props.blockId);
             if (typeof updateData === "string") {
               // received a url
               updateData = {
@@ -60,7 +62,7 @@ export const UploadTab = <
                 },
               };
             }
-            editor.updateBlock(block, updateData);
+            editor.updateBlock(props.blockId, updateData);
           } catch (e) {
             setUploadFailed(true);
           } finally {
@@ -71,7 +73,7 @@ export const UploadTab = <
 
       upload(file);
     },
-    [block, editor, setLoading],
+    [props.blockId, editor, setLoading],
   );
 
   const spec = editor.schema.blockSpecs[block.type];
diff --git a/packages/react/src/components/FilePanel/FilePanel.tsx b/packages/react/src/components/FilePanel/FilePanel.tsx
index 270e749dcd..9365571025 100644
--- a/packages/react/src/components/FilePanel/FilePanel.tsx
+++ b/packages/react/src/components/FilePanel/FilePanel.tsx
@@ -31,8 +31,7 @@ export const FilePanel = <
   I extends InlineContentSchema = DefaultInlineContentSchema,
   S extends StyleSchema = DefaultStyleSchema,
 >(
-  props: FilePanelProps<I, S> &
-    Partial<Pick<PanelProps, "defaultOpenTab" | "tabs">>,
+  props: FilePanelProps & Partial<Pick<PanelProps, "defaultOpenTab" | "tabs">>,
 ) => {
   const Components = useComponentsContext()!;
   const dict = useDictionary();
@@ -46,13 +45,15 @@ export const FilePanel = <
       ? [
           {
             name: dict.file_panel.upload.title,
-            tabPanel: <UploadTab block={props.block} setLoading={setLoading} />,
+            tabPanel: (
+              <UploadTab blockId={props.blockId} setLoading={setLoading} />
+            ),
           },
         ]
       : []),
     {
       name: dict.file_panel.embed.title,
-      tabPanel: <EmbedTab block={props.block} />,
+      tabPanel: <EmbedTab blockId={props.blockId} />,
     },
   ];
 
diff --git a/packages/react/src/components/FilePanel/FilePanelController.tsx b/packages/react/src/components/FilePanel/FilePanelController.tsx
index 7fb98b2bab..04e611c1f4 100644
--- a/packages/react/src/components/FilePanel/FilePanelController.tsx
+++ b/packages/react/src/components/FilePanel/FilePanelController.tsx
@@ -1,68 +1,53 @@
-import {
-  BlockSchema,
-  DefaultBlockSchema,
-  DefaultInlineContentSchema,
-  DefaultStyleSchema,
-  InlineContentSchema,
-  StyleSchema,
-} from "@blocknote/core";
-import { UseFloatingOptions, flip, offset } from "@floating-ui/react";
-import { FC } from "react";
+import { FilePanel as FilePanelExtension } from "@blocknote/core/extensions";
+import { flip, offset } from "@floating-ui/react";
+import { FC, useEffect, useMemo, useState } from "react";
 
-import { useBlockNoteEditor } from "../../hooks/useBlockNoteEditor.js";
-import { useUIElementPositioning } from "../../hooks/useUIElementPositioning.js";
-import { useUIPluginState } from "../../hooks/useUIPluginState.js";
 import { FilePanel } from "./FilePanel.js";
 import { FilePanelProps } from "./FilePanelProps.js";
+import { BlockPopover } from "../Popovers/BlockPopover.js";
+import { FloatingUIOptions } from "../Popovers/FloatingUIOptions.js";
+import { useExtensionState } from "../../hooks/useExtension.js";
 
-export const FilePanelController = <
-  B extends BlockSchema = DefaultBlockSchema,
-  I extends InlineContentSchema = DefaultInlineContentSchema,
-  S extends StyleSchema = DefaultStyleSchema,
->(props: {
-  filePanel?: FC<FilePanelProps<I, S>>;
-  floatingOptions?: Partial<UseFloatingOptions>;
+export const FilePanelController = (props: {
+  filePanel?: FC<FilePanelProps>;
+  floatingUIOptions?: FloatingUIOptions;
 }) => {
-  const editor = useBlockNoteEditor<B, I, S>();
-
-  if (!editor.filePanel) {
-    throw new Error(
-      "FileToolbarController can only be used when BlockNote editor schema contains file block",
-    );
-  }
-
-  const state = useUIPluginState(
-    editor.filePanel.onUpdate.bind(editor.filePanel),
-  );
-
-  const { isMounted, ref, style, getFloatingProps } = useUIElementPositioning(
-    state?.show || false,
-    state?.referencePos || null,
-    5000,
-    {
-      placement: "bottom",
-      middleware: [offset(10), flip()],
-      onOpenChange: (open) => {
-        if (!open) {
-          editor.filePanel!.closeMenu();
-          editor.focus();
-        }
+  const [open, setOpen] = useState(false);
+
+  const blockId = useExtensionState(FilePanelExtension, {
+    selector: (state) => state.blockId,
+  });
+  useEffect(() => {
+    setOpen(blockId !== undefined);
+  }, [blockId]);
+
+  const floatingUIOptions = useMemo<FloatingUIOptions>(
+    () => ({
+      useFloatingOptions: {
+        open,
+        // Needed as hooks like `useDismiss` call `onOpenChange` to change the
+        // open state.
+        onOpenChange: setOpen,
+        middleware: [offset(10), flip()],
       },
-      ...props.floatingOptions,
-    },
+      useDismissProps: {
+        outsidePress: false,
+      },
+      elementProps: {
+        style: {
+          zIndex: 90,
+        },
+      },
+      ...props.floatingUIOptions,
+    }),
+    [open, props.floatingUIOptions],
   );
 
-  if (!isMounted || !state) {
-    return null;
-  }
-
-  const { show, referencePos, ...data } = state;
-
   const Component = props.filePanel || FilePanel;
 
   return (
-    <div ref={ref} style={style} {...getFloatingProps()}>
-      <Component {...data} />
-    </div>
+    <BlockPopover blockId={blockId} {...floatingUIOptions}>
+      {blockId && <Component blockId={blockId} />}
+    </BlockPopover>
   );
 };
diff --git a/packages/react/src/components/FilePanel/FilePanelProps.ts b/packages/react/src/components/FilePanel/FilePanelProps.ts
index ac1420e011..b953fb75f1 100644
--- a/packages/react/src/components/FilePanel/FilePanelProps.ts
+++ b/packages/react/src/components/FilePanel/FilePanelProps.ts
@@ -1,13 +1,3 @@
-import {
-  DefaultInlineContentSchema,
-  DefaultStyleSchema,
-  FilePanelState,
-  InlineContentSchema,
-  StyleSchema,
-  UiElementPosition,
-} from "@blocknote/core";
-
-export type FilePanelProps<
-  I extends InlineContentSchema = DefaultInlineContentSchema,
-  S extends StyleSchema = DefaultStyleSchema,
-> = Omit<FilePanelState<I, S>, keyof UiElementPosition>;
+export type FilePanelProps = {
+  blockId: string;
+};
diff --git a/packages/react/src/components/FormattingToolbar/DefaultButtons/AddCommentButton.tsx b/packages/react/src/components/FormattingToolbar/DefaultButtons/AddCommentButton.tsx
index 3e85b2994f..b65a412486 100644
--- a/packages/react/src/components/FormattingToolbar/DefaultButtons/AddCommentButton.tsx
+++ b/packages/react/src/components/FormattingToolbar/DefaultButtons/AddCommentButton.tsx
@@ -1,33 +1,23 @@
-import { BlockSchema, InlineContentSchema, StyleSchema } from "@blocknote/core";
+import { Comments, FormattingToolbar } from "@blocknote/core/extensions";
 import { useCallback } from "react";
 import { RiChat3Line } from "react-icons/ri";
 
 import { useComponentsContext } from "../../../editor/ComponentsContext.js";
 import { useBlockNoteEditor } from "../../../hooks/useBlockNoteEditor.js";
 import { useDictionary } from "../../../i18n/dictionary.js";
+import { useExtension } from "../../../hooks/useExtension.js";
 
-export const AddCommentButton = () => {
+export const AddCommentButtonInner = () => {
   const dict = useDictionary();
   const Components = useComponentsContext()!;
 
-  const editor = useBlockNoteEditor<
-    BlockSchema,
-    InlineContentSchema,
-    StyleSchema
-  >();
+  const comments = useExtension(Comments);
+  const { store } = useExtension(FormattingToolbar);
 
   const onClick = useCallback(() => {
-    editor.comments?.startPendingComment();
-    editor.formattingToolbar.closeMenu();
-  }, [editor]);
-
-  if (
-    // We manually check if a comment extension (like liveblocks) is installed
-    // By adding default support for this, the user doesn't need to customize the formatting toolbar
-    !editor.comments
-  ) {
-    return null;
-  }
+    comments.startPendingComment();
+    store.setState(false);
+  }, [comments, store]);
 
   return (
     <Components.FormattingToolbar.Button
@@ -39,3 +29,13 @@ export const AddCommentButton = () => {
     />
   );
 };
+
+export const AddCommentButton = () => {
+  const editor = useBlockNoteEditor<any, any, any>();
+
+  if (!editor.getExtension(Comments)) {
+    return null;
+  }
+
+  return <AddCommentButtonInner />;
+};
diff --git a/packages/react/src/components/FormattingToolbar/DefaultButtons/BasicTextStyleButton.tsx b/packages/react/src/components/FormattingToolbar/DefaultButtons/BasicTextStyleButton.tsx
index 8311c79ce0..f2ae1e9ac6 100644
--- a/packages/react/src/components/FormattingToolbar/DefaultButtons/BasicTextStyleButton.tsx
+++ b/packages/react/src/components/FormattingToolbar/DefaultButtons/BasicTextStyleButton.tsx
@@ -5,7 +5,7 @@ import {
   StyleSchema,
   formatKeyboardShortcut,
 } from "@blocknote/core";
-import { useMemo, useState } from "react";
+import { useCallback } from "react";
 import { IconType } from "react-icons";
 import {
   RiBold,
@@ -17,8 +17,7 @@ import {
 
 import { useComponentsContext } from "../../../editor/ComponentsContext.js";
 import { useBlockNoteEditor } from "../../../hooks/useBlockNoteEditor.js";
-import { useEditorContentOrSelectionChange } from "../../../hooks/useEditorContentOrSelectionChange.js";
-import { useSelectedBlocks } from "../../../hooks/useSelectedBlocks.js";
+import { useEditorState } from "../../../hooks/useEditorState.js";
 import { useDictionary } from "../../../i18n/dictionary.js";
 
 type BasicTextStyle = "bold" | "italic" | "underline" | "strike" | "code";
@@ -63,45 +62,40 @@ export const BasicTextStyleButton = <Style extends BasicTextStyle>(props: {
     StyleSchema
   >();
 
-  const basicTextStyleInSchema = checkBasicTextStyleInSchema(
-    props.basicTextStyle,
+  const state = useEditorState({
     editor,
+    selector: ({ editor }) => {
+      // Do not show if:
+      if (
+        // The editor is read-only.
+        !editor.isEditable ||
+        // The style is not in the schema.
+        !checkBasicTextStyleInSchema(props.basicTextStyle, editor) ||
+        // None of the selected blocks have inline content
+        !(
+          editor.getSelection()?.blocks || [
+            editor.getTextCursorPosition().block,
+          ]
+        ).find((block) => block.content !== undefined)
+      ) {
+        return undefined;
+      }
+
+      return props.basicTextStyle in editor.getActiveStyles()
+        ? { active: true }
+        : { active: false };
+    },
+  });
+
+  const toggleStyle = useCallback(
+    (style: typeof props.basicTextStyle) => {
+      editor.focus();
+      editor.toggleStyles({ [style]: true } as any);
+    },
+    [editor, props],
   );
 
-  const selectedBlocks = useSelectedBlocks(editor);
-
-  const [active, setActive] = useState<boolean>(
-    props.basicTextStyle in editor.getActiveStyles(),
-  );
-
-  useEditorContentOrSelectionChange(() => {
-    if (basicTextStyleInSchema) {
-      setActive(props.basicTextStyle in editor.getActiveStyles());
-    }
-  }, editor);
-
-  const toggleStyle = (style: typeof props.basicTextStyle) => {
-    editor.focus();
-
-    if (!basicTextStyleInSchema) {
-      return;
-    }
-
-    if (editor.schema.styleSchema[style].propSchema !== "boolean") {
-      throw new Error("can only toggle boolean styles");
-    }
-    editor.toggleStyles({ [style]: true } as any);
-  };
-
-  const show = useMemo(() => {
-    if (!basicTextStyleInSchema) {
-      return false;
-    }
-    // Also don't show when none of the selected blocks have text content
-    return !!selectedBlocks.find((block) => block.content !== undefined);
-  }, [basicTextStyleInSchema, selectedBlocks]);
-
-  if (!show || !editor.isEditable) {
+  if (state === undefined) {
     return null;
   }
 
@@ -111,7 +105,7 @@ export const BasicTextStyleButton = <Style extends BasicTextStyle>(props: {
       className="bn-button"
       data-test={props.basicTextStyle}
       onClick={() => toggleStyle(props.basicTextStyle)}
-      isSelected={active}
+      isSelected={state.active}
       label={dict.formatting_toolbar[props.basicTextStyle].tooltip}
       mainTooltip={dict.formatting_toolbar[props.basicTextStyle].tooltip}
       secondaryTooltip={formatKeyboardShortcut(
diff --git a/packages/react/src/components/FormattingToolbar/DefaultButtons/ColorStyleButton.tsx b/packages/react/src/components/FormattingToolbar/DefaultButtons/ColorStyleButton.tsx
index 1a9ed2da07..bb978b8a6c 100644
--- a/packages/react/src/components/FormattingToolbar/DefaultButtons/ColorStyleButton.tsx
+++ b/packages/react/src/components/FormattingToolbar/DefaultButtons/ColorStyleButton.tsx
@@ -4,12 +4,11 @@ import {
   InlineContentSchema,
   StyleSchema,
 } from "@blocknote/core";
-import { useCallback, useMemo, useState } from "react";
+import { useCallback } from "react";
 
 import { useComponentsContext } from "../../../editor/ComponentsContext.js";
 import { useBlockNoteEditor } from "../../../hooks/useBlockNoteEditor.js";
-import { useEditorContentOrSelectionChange } from "../../../hooks/useEditorContentOrSelectionChange.js";
-import { useSelectedBlocks } from "../../../hooks/useSelectedBlocks.js";
+import { useEditorState } from "../../../hooks/useEditorState.js";
 import { useDictionary } from "../../../i18n/dictionary.js";
 import { ColorIcon } from "../../ColorPicker/ColorIcon.js";
 import { ColorPicker } from "../../ColorPicker/ColorPicker.js";
@@ -53,29 +52,40 @@ export const ColorStyleButton = () => {
   const textColorInSchema = checkColorInSchema("text", editor);
   const backgroundColorInSchema = checkColorInSchema("background", editor);
 
-  const selectedBlocks = useSelectedBlocks(editor);
+  const state = useEditorState({
+    editor,
+    selector: ({ editor }) => {
+      // Do not show if:
+      if (
+        // The editor is read-only.
+        !editor.isEditable ||
+        // None of the selected blocks have inline content
+        !(
+          editor.getSelection()?.blocks || [
+            editor.getTextCursorPosition().block,
+          ]
+        ).find((block) => block.content !== undefined)
+      ) {
+        return undefined;
+      }
 
-  const [currentTextColor, setCurrentTextColor] = useState<string>(
-    textColorInSchema
-      ? editor.getActiveStyles().textColor || "default"
-      : "default",
-  );
-  const [currentBackgroundColor, setCurrentBackgroundColor] = useState<string>(
-    backgroundColorInSchema
-      ? editor.getActiveStyles().backgroundColor || "default"
-      : "default",
-  );
+      const textColorInSchema = checkColorInSchema("text", editor);
+      const backgroundColorInSchema = checkColorInSchema("background", editor);
+
+      if (!textColorInSchema && !backgroundColorInSchema) {
+        return undefined;
+      }
 
-  useEditorContentOrSelectionChange(() => {
-    if (textColorInSchema) {
-      setCurrentTextColor(editor.getActiveStyles().textColor || "default");
-    }
-    if (backgroundColorInSchema) {
-      setCurrentBackgroundColor(
-        editor.getActiveStyles().backgroundColor || "default",
-      );
-    }
-  }, editor);
+      return {
+        textColor: (textColorInSchema
+          ? editor.getActiveStyles().textColor || "default"
+          : undefined) as string | undefined,
+        backgroundColor: (backgroundColorInSchema
+          ? editor.getActiveStyles().backgroundColor || "default"
+          : undefined) as string | undefined,
+      };
+    },
+  });
 
   const setTextColor = useCallback(
     (color: string) => {
@@ -117,21 +127,7 @@ export const ColorStyleButton = () => {
     [backgroundColorInSchema, editor],
   );
 
-  const show = useMemo(() => {
-    if (!textColorInSchema && !backgroundColorInSchema) {
-      return false;
-    }
-
-    for (const block of selectedBlocks) {
-      if (block.content !== undefined) {
-        return true;
-      }
-    }
-
-    return false;
-  }, [backgroundColorInSchema, selectedBlocks, textColorInSchema]);
-
-  if (!show || !editor.isEditable) {
+  if (state === undefined) {
     return null;
   }
 
@@ -145,8 +141,8 @@ export const ColorStyleButton = () => {
           mainTooltip={dict.formatting_toolbar.colors.tooltip}
           icon={
             <ColorIcon
-              textColor={currentTextColor}
-              backgroundColor={currentBackgroundColor}
+              textColor={state.textColor}
+              backgroundColor={state.backgroundColor}
               size={20}
             />
           }
@@ -157,17 +153,17 @@ export const ColorStyleButton = () => {
       >
         <ColorPicker
           text={
-            textColorInSchema
+            state.textColor
               ? {
-                  color: currentTextColor,
+                  color: state.textColor,
                   setColor: setTextColor,
                 }
               : undefined
           }
           background={
-            backgroundColorInSchema
+            state.backgroundColor
               ? {
-                  color: currentBackgroundColor,
+                  color: state.backgroundColor,
                   setColor: setBackgroundColor,
                 }
               : undefined
diff --git a/packages/react/src/components/FormattingToolbar/DefaultButtons/CreateLinkButton.tsx b/packages/react/src/components/FormattingToolbar/DefaultButtons/CreateLinkButton.tsx
index f5062f9309..4c0bdcc5ae 100644
--- a/packages/react/src/components/FormattingToolbar/DefaultButtons/CreateLinkButton.tsx
+++ b/packages/react/src/components/FormattingToolbar/DefaultButtons/CreateLinkButton.tsx
@@ -1,4 +1,4 @@
-import { useCallback, useEffect, useMemo, useState } from "react";
+import { useEffect, useState } from "react";
 import { RiLink } from "react-icons/ri";
 
 import {
@@ -12,8 +12,7 @@ import {
 
 import { useComponentsContext } from "../../../editor/ComponentsContext.js";
 import { useBlockNoteEditor } from "../../../hooks/useBlockNoteEditor.js";
-import { useEditorContentOrSelectionChange } from "../../../hooks/useEditorContentOrSelectionChange.js";
-import { useSelectedBlocks } from "../../../hooks/useSelectedBlocks.js";
+import { useEditorState } from "../../../hooks/useEditorState.js";
 import { useDictionary } from "../../../i18n/dictionary.js";
 import { EditLinkMenuItems } from "../../LinkToolbar/EditLinkMenuItems.js";
 
@@ -45,31 +44,49 @@ export const CreateLinkButton = () => {
   const Components = useComponentsContext()!;
   const dict = useDictionary();
 
-  const linkInSchema = checkLinkInSchema(editor);
-
-  const selectedBlocks = useSelectedBlocks(editor);
-
-  const [opened, setOpened] = useState(false);
-  const [url, setUrl] = useState<string>(editor.getSelectedLinkUrl() || "");
-  const [text, setText] = useState<string>(editor.getSelectedText());
+  const [showPopover, setShowPopover] = useState(false);
+
+  const state = useEditorState({
+    editor,
+    selector: ({ editor }) => {
+      // Do not show if:
+      if (
+        // The editor is read-only.
+        !editor.isEditable ||
+        // Links are not in the schema.
+        !checkLinkInSchema(editor) ||
+        // Table cells are selected.
+        isTableCellSelection(editor.prosemirrorState.selection) ||
+        // None of the selected blocks have inline content
+        !(
+          editor.getSelection()?.blocks || [
+            editor.getTextCursorPosition().block,
+          ]
+        ).find((block) => block.content !== undefined)
+      ) {
+        return undefined;
+      }
 
-  useEditorContentOrSelectionChange(() => {
-    setText(editor.getSelectedText() || "");
-    setUrl(editor.getSelectedLinkUrl() || "");
-  }, editor);
+      return {
+        url: editor.getSelectedLinkUrl(),
+        text: editor.getSelectedText(),
+        range: {
+          from: editor.prosemirrorState.selection.from,
+          to: editor.prosemirrorState.selection.to,
+        },
+      };
+    },
+  });
 
+  // Makes Ctrl+K/Meta+K open link creation popover.
   useEffect(() => {
     const callback = (event: KeyboardEvent) => {
       if ((event.ctrlKey || event.metaKey) && event.key === "k") {
-        setOpened(true);
+        setShowPopover(true);
         event.preventDefault();
       }
     };
 
-    if (editor.headless) {
-      return;
-    }
-
     editor.prosemirrorView.dom.addEventListener("keydown", callback);
 
     return () => {
@@ -77,43 +94,12 @@ export const CreateLinkButton = () => {
     };
   }, [editor.prosemirrorView, editor.headless]);
 
-  const update = useCallback(
-    (url: string) => {
-      editor.createLink(url);
-      setOpened(false);
-      editor.focus();
-    },
-    [editor],
-  );
-
-  const isTableSelection = editor.transact((tr) =>
-    isTableCellSelection(tr.selection),
-  );
-
-  const show = useMemo(() => {
-    if (!linkInSchema) {
-      return false;
-    }
-
-    for (const block of selectedBlocks) {
-      if (block.content === undefined) {
-        return false;
-      }
-    }
-
-    return !isTableSelection;
-  }, [linkInSchema, selectedBlocks, isTableSelection]);
-
-  if (
-    !show ||
-    !("link" in editor.schema.inlineContentSchema) ||
-    !editor.isEditable
-  ) {
+  if (state === undefined) {
     return null;
   }
 
   return (
-    <Components.Generic.Popover.Root opened={opened}>
+    <Components.Generic.Popover.Root opened={showPopover}>
       <Components.Generic.Popover.Trigger>
         {/* TODO: hide tooltip on click */}
         <Components.FormattingToolbar.Button
@@ -126,7 +112,7 @@ export const CreateLinkButton = () => {
             dict.generic.ctrl_shortcut,
           )}
           icon={<RiLink />}
-          onClick={() => setOpened(true)}
+          onClick={() => setShowPopover(true)}
         />
       </Components.Generic.Popover.Trigger>
       <Components.Generic.Popover.Content
@@ -134,9 +120,9 @@ export const CreateLinkButton = () => {
         variant={"form-popover"}
       >
         <EditLinkMenuItems
-          url={url}
-          text={text}
-          editLink={update}
+          url={state.url || ""}
+          text={state.text}
+          range={state.range}
           showTextField={false}
         />
       </Components.Generic.Popover.Content>
diff --git a/packages/react/src/components/FormattingToolbar/DefaultButtons/FileCaptionButton.tsx b/packages/react/src/components/FormattingToolbar/DefaultButtons/FileCaptionButton.tsx
index 38b1fef90c..72200ba723 100644
--- a/packages/react/src/components/FormattingToolbar/DefaultButtons/FileCaptionButton.tsx
+++ b/packages/react/src/components/FormattingToolbar/DefaultButtons/FileCaptionButton.tsx
@@ -9,14 +9,14 @@ import {
   ChangeEvent,
   KeyboardEvent,
   useCallback,
-  useMemo,
+  useEffect,
   useState,
 } from "react";
 import { RiInputField } from "react-icons/ri";
 
 import { useComponentsContext } from "../../../editor/ComponentsContext.js";
 import { useBlockNoteEditor } from "../../../hooks/useBlockNoteEditor.js";
-import { useSelectedBlocks } from "../../../hooks/useSelectedBlocks.js";
+import { useEditorState } from "../../../hooks/useEditorState.js";
 import { useDictionary } from "../../../i18n/dictionary.js";
 
 export const FileCaptionButton = () => {
@@ -29,58 +29,76 @@ export const FileCaptionButton = () => {
     StyleSchema
   >();
 
-  const [currentEditingCaption, setCurrentEditingCaption] = useState<string>();
+  const state = useEditorState({
+    editor,
+    selector: ({ editor }) => {
+      if (!editor.isEditable) {
+        return undefined;
+      }
 
-  const selectedBlocks = useSelectedBlocks(editor);
+      const selectedBlocks = editor.getSelection()?.blocks || [
+        editor.getTextCursorPosition().block,
+      ];
 
-  const fileBlock = useMemo(() => {
-    // Checks if only one block is selected.
-    if (selectedBlocks.length !== 1) {
-      return undefined;
-    }
+      if (selectedBlocks.length !== 1) {
+        return undefined;
+      }
+
+      const block = selectedBlocks[0];
+
+      if (
+        !blockHasType(block, editor, block.type, {
+          url: "string",
+          caption: "string",
+        })
+      ) {
+        return undefined;
+      }
+
+      return {
+        blockId: block.id,
+        blockType: block.type,
+        caption: block.props.caption,
+      };
+    },
+  });
 
-    const block = selectedBlocks[0];
+  const [currentEditingCaption, setCurrentEditingCaption] = useState<string>();
 
-    if (
-      blockHasType(block, editor, block.type, {
-        url: "string",
-        caption: "string",
-      })
-    ) {
-      setCurrentEditingCaption(block.props.caption);
-      return block;
+  useEffect(() => {
+    if (!state) {
+      return;
     }
+    setCurrentEditingCaption(state.caption);
+  }, [state]);
 
-    return undefined;
-  }, [editor, selectedBlocks]);
+  const handleChange = useCallback(
+    (event: ChangeEvent<HTMLInputElement>) =>
+      setCurrentEditingCaption(event.currentTarget.value),
+    [],
+  );
 
   const handleEnter = useCallback(
     (event: KeyboardEvent) => {
       if (
-        fileBlock &&
-        editorHasBlockWithType(editor, fileBlock.type, {
+        state !== undefined &&
+        editorHasBlockWithType(editor, state.blockType, {
           caption: "string",
         }) &&
         event.key === "Enter"
       ) {
         event.preventDefault();
-        editor.updateBlock(fileBlock, {
+        editor.updateBlock(state.blockId, {
           props: {
             caption: currentEditingCaption,
           },
         });
       }
     },
-    [currentEditingCaption, editor, fileBlock],
-  );
-
-  const handleChange = useCallback(
-    (event: ChangeEvent<HTMLInputElement>) =>
-      setCurrentEditingCaption(event.currentTarget.value),
-    [],
+    [currentEditingCaption, editor, state],
   );
 
-  if (!fileBlock || fileBlock.props.url === "" || !editor.isEditable) {
+  if (state === undefined) {
     return null;
   }
 
@@ -92,7 +110,6 @@ export const FileCaptionButton = () => {
           label={dict.formatting_toolbar.file_caption.tooltip}
           mainTooltip={dict.formatting_toolbar.file_caption.tooltip}
           icon={<RiInputField />}
-          isSelected={fileBlock.props.caption !== ""}
         />
       </Components.Generic.Popover.Trigger>
       <Components.Generic.Popover.Content
diff --git a/packages/react/src/components/FormattingToolbar/DefaultButtons/FileDeleteButton.tsx b/packages/react/src/components/FormattingToolbar/DefaultButtons/FileDeleteButton.tsx
index 2ea096de36..3e46b4163b 100644
--- a/packages/react/src/components/FormattingToolbar/DefaultButtons/FileDeleteButton.tsx
+++ b/packages/react/src/components/FormattingToolbar/DefaultButtons/FileDeleteButton.tsx
@@ -4,12 +4,12 @@ import {
   InlineContentSchema,
   StyleSchema,
 } from "@blocknote/core";
-import { useCallback, useMemo } from "react";
+import { useCallback } from "react";
 import { RiDeleteBin7Line } from "react-icons/ri";
 
 import { useComponentsContext } from "../../../editor/ComponentsContext.js";
 import { useBlockNoteEditor } from "../../../hooks/useBlockNoteEditor.js";
-import { useSelectedBlocks } from "../../../hooks/useSelectedBlocks.js";
+import { useEditorState } from "../../../hooks/useEditorState.js";
 import { useDictionary } from "../../../i18n/dictionary.js";
 
 export const FileDeleteButton = () => {
@@ -22,29 +22,43 @@ export const FileDeleteButton = () => {
     StyleSchema
   >();
 
-  const selectedBlocks = useSelectedBlocks(editor);
+  const state = useEditorState({
+    editor,
+    selector: ({ editor }) => {
+      if (!editor.isEditable) {
+        return undefined;
+      }
 
-  const fileBlock = useMemo(() => {
-    // Checks if only one block is selected.
-    if (selectedBlocks.length !== 1) {
-      return undefined;
-    }
+      const selectedBlocks = editor.getSelection()?.blocks || [
+        editor.getTextCursorPosition().block,
+      ];
 
-    const block = selectedBlocks[0];
+      if (selectedBlocks.length !== 1) {
+        return undefined;
+      }
 
-    if (blockHasType(block, editor, block.type, { url: "string" })) {
-      return block;
-    }
+      const block = selectedBlocks[0];
 
-    return undefined;
-  }, [editor, selectedBlocks]);
+      if (
+        !blockHasType(block, editor, block.type, {
+          url: "string",
+        })
+      ) {
+        return undefined;
+      }
+
+      return { blockId: block.id, blockType: block.type };
+    },
+  });
 
   const onClick = useCallback(() => {
-    editor.focus();
-    editor.removeBlocks([fileBlock!]);
-  }, [editor, fileBlock]);
+    if (state !== undefined) {
+      editor.focus();
+      editor.removeBlocks([state.blockId]);
+    }
+  }, [editor, state]);
 
-  if (!fileBlock || fileBlock.props.url === "" || !editor.isEditable) {
+  if (state === undefined) {
     return null;
   }
 
@@ -52,11 +66,11 @@ export const FileDeleteButton = () => {
     <Components.FormattingToolbar.Button
       className={"bn-button"}
       label={
-        dict.formatting_toolbar.file_delete.tooltip[fileBlock.type] ||
+        dict.formatting_toolbar.file_delete.tooltip[state.blockType] ||
         dict.formatting_toolbar.file_delete.tooltip["file"]
       }
       mainTooltip={
-        dict.formatting_toolbar.file_delete.tooltip[fileBlock.type] ||
+        dict.formatting_toolbar.file_delete.tooltip[state.blockType] ||
         dict.formatting_toolbar.file_delete.tooltip["file"]
       }
       icon={<RiDeleteBin7Line />}
diff --git a/packages/react/src/components/FormattingToolbar/DefaultButtons/FileDownloadButton.tsx b/packages/react/src/components/FormattingToolbar/DefaultButtons/FileDownloadButton.tsx
index d19e202564..8542f880fa 100644
--- a/packages/react/src/components/FormattingToolbar/DefaultButtons/FileDownloadButton.tsx
+++ b/packages/react/src/components/FormattingToolbar/DefaultButtons/FileDownloadButton.tsx
@@ -4,12 +4,12 @@ import {
   InlineContentSchema,
   StyleSchema,
 } from "@blocknote/core";
-import { useCallback, useMemo } from "react";
+import { useCallback } from "react";
 import { RiDownload2Fill } from "react-icons/ri";
 
 import { useComponentsContext } from "../../../editor/ComponentsContext.js";
 import { useBlockNoteEditor } from "../../../hooks/useBlockNoteEditor.js";
-import { useSelectedBlocks } from "../../../hooks/useSelectedBlocks.js";
+import { useEditorState } from "../../../hooks/useEditorState.js";
 import { useDictionary } from "../../../i18n/dictionary.js";
 import { sanitizeUrl } from "../../../util/sanitizeUrl.js";
 
@@ -23,40 +23,52 @@ export const FileDownloadButton = () => {
     StyleSchema
   >();
 
-  const selectedBlocks = useSelectedBlocks(editor);
+  const state = useEditorState({
+    editor,
+    selector: ({ editor }) => {
+      if (!editor.isEditable) {
+        return undefined;
+      }
 
-  const fileBlock = useMemo(() => {
-    // Checks if only one block is selected.
-    if (selectedBlocks.length !== 1) {
-      return undefined;
-    }
+      const selectedBlocks = editor.getSelection()?.blocks || [
+        editor.getTextCursorPosition().block,
+      ];
+
+      if (selectedBlocks.length !== 1) {
+        return undefined;
+      }
 
-    const block = selectedBlocks[0];
+      const block = selectedBlocks[0];
 
-    if (blockHasType(block, editor, block.type, { url: "string" })) {
-      return block;
-    }
+      if (
+        !blockHasType(block, editor, block.type, {
+          url: "string",
+        })
+      ) {
+        return undefined;
+      }
 
-    return undefined;
-  }, [editor, selectedBlocks]);
+      return { blockId: block.id, blockType: block.type, url: block.props.url };
+    },
+  });
 
   const onClick = useCallback(() => {
-    if (fileBlock && fileBlock.props.url) {
+    if (state !== undefined) {
       editor.focus();
 
       if (!editor.resolveFileUrl) {
-        window.open(sanitizeUrl(fileBlock.props.url, window.location.href));
+        window.open(sanitizeUrl(state.url, window.location.href));
       } else {
         editor
-          .resolveFileUrl(fileBlock.props.url)
+          .resolveFileUrl(state.url)
           .then((downloadUrl) =>
             window.open(sanitizeUrl(downloadUrl, window.location.href)),
           );
       }
     }
-  }, [editor, fileBlock]);
+  }, [editor, state]);
 
-  if (!fileBlock || fileBlock.props.url === "") {
+  if (state === undefined) {
     return null;
   }
 
@@ -64,11 +76,11 @@ export const FileDownloadButton = () => {
     <Components.FormattingToolbar.Button
       className={"bn-button"}
       label={
-        dict.formatting_toolbar.file_download.tooltip[fileBlock.type] ||
+        dict.formatting_toolbar.file_download.tooltip[state.blockType] ||
         dict.formatting_toolbar.file_download.tooltip["file"]
       }
       mainTooltip={
-        dict.formatting_toolbar.file_download.tooltip[fileBlock.type] ||
+        dict.formatting_toolbar.file_download.tooltip[state.blockType] ||
         dict.formatting_toolbar.file_download.tooltip["file"]
       }
       icon={<RiDownload2Fill />}
diff --git a/packages/react/src/components/FormattingToolbar/DefaultButtons/FilePreviewButton.tsx b/packages/react/src/components/FormattingToolbar/DefaultButtons/FilePreviewButton.tsx
index 9e8fcb0425..c770f91c4a 100644
--- a/packages/react/src/components/FormattingToolbar/DefaultButtons/FilePreviewButton.tsx
+++ b/packages/react/src/components/FormattingToolbar/DefaultButtons/FilePreviewButton.tsx
@@ -5,12 +5,12 @@ import {
   InlineContentSchema,
   StyleSchema,
 } from "@blocknote/core";
-import { useCallback, useMemo } from "react";
+import { useCallback } from "react";
 import { RiImageAddFill } from "react-icons/ri";
 
 import { useComponentsContext } from "../../../editor/ComponentsContext.js";
 import { useBlockNoteEditor } from "../../../hooks/useBlockNoteEditor.js";
-import { useSelectedBlocks } from "../../../hooks/useSelectedBlocks.js";
+import { useEditorState } from "../../../hooks/useEditorState.js";
 import { useDictionary } from "../../../i18n/dictionary.js";
 
 export const FilePreviewButton = () => {
@@ -23,44 +23,56 @@ export const FilePreviewButton = () => {
     StyleSchema
   >();
 
-  const selectedBlocks = useSelectedBlocks(editor);
+  const state = useEditorState({
+    editor,
+    selector: ({ editor }) => {
+      if (!editor.isEditable) {
+        return undefined;
+      }
 
-  const fileBlock = useMemo(() => {
-    // Checks if only one block is selected.
-    if (selectedBlocks.length !== 1) {
-      return undefined;
-    }
+      const selectedBlocks = editor.getSelection()?.blocks || [
+        editor.getTextCursorPosition().block,
+      ];
 
-    const block = selectedBlocks[0];
+      if (selectedBlocks.length !== 1) {
+        return undefined;
+      }
 
-    if (
-      blockHasType(block, editor, block.type, {
-        url: "string",
-        showPreview: "boolean",
-      })
-    ) {
-      return block;
-    }
+      const block = selectedBlocks[0];
+
+      if (
+        !blockHasType(block, editor, block.type, {
+          url: "string",
+          showPreview: "boolean",
+        })
+      ) {
+        return undefined;
+      }
 
-    return undefined;
-  }, [editor, selectedBlocks]);
+      return {
+        blockId: block.id,
+        blockType: block.type,
+        showPreview: block.props.showPreview,
+      };
+    },
+  });
 
   const onClick = useCallback(() => {
     if (
-      fileBlock &&
-      editorHasBlockWithType(editor, fileBlock.type, {
+      state !== undefined &&
+      editorHasBlockWithType(editor, state.blockType, {
         showPreview: "boolean",
       })
     ) {
-      editor.updateBlock(fileBlock, {
+      editor.updateBlock(state.blockId, {
         props: {
-          showPreview: !fileBlock.props.showPreview,
+          showPreview: !state.showPreview,
         },
       });
     }
-  }, [editor, fileBlock]);
+  }, [editor, state]);
 
-  if (!fileBlock || fileBlock.props.url === "" || !editor.isEditable) {
+  if (state === undefined) {
     return null;
   }
 
@@ -70,7 +82,7 @@ export const FilePreviewButton = () => {
       label={"Toggle preview"}
       mainTooltip={dict.formatting_toolbar.file_preview_toggle.tooltip}
       icon={<RiImageAddFill />}
-      isSelected={fileBlock.props.showPreview}
+      isSelected={state.showPreview}
       onClick={onClick}
     />
   );
diff --git a/packages/react/src/components/FormattingToolbar/DefaultButtons/FileRenameButton.tsx b/packages/react/src/components/FormattingToolbar/DefaultButtons/FileRenameButton.tsx
index 583494917f..829ad34391 100644
--- a/packages/react/src/components/FormattingToolbar/DefaultButtons/FileRenameButton.tsx
+++ b/packages/react/src/components/FormattingToolbar/DefaultButtons/FileRenameButton.tsx
@@ -9,14 +9,14 @@ import {
   ChangeEvent,
   KeyboardEvent,
   useCallback,
-  useMemo,
+  useEffect,
   useState,
 } from "react";
 import { RiFontFamily } from "react-icons/ri";
 
 import { useComponentsContext } from "../../../editor/ComponentsContext.js";
 import { useBlockNoteEditor } from "../../../hooks/useBlockNoteEditor.js";
-import { useSelectedBlocks } from "../../../hooks/useSelectedBlocks.js";
+import { useEditorState } from "../../../hooks/useEditorState.js";
 import { useDictionary } from "../../../i18n/dictionary.js";
 
 export const FileRenameButton = () => {
@@ -29,58 +29,76 @@ export const FileRenameButton = () => {
     StyleSchema
   >();
 
-  const [currentEditingName, setCurrentEditingName] = useState<string>();
+  const state = useEditorState({
+    editor,
+    selector: ({ editor }) => {
+      if (!editor.isEditable) {
+        return undefined;
+      }
 
-  const selectedBlocks = useSelectedBlocks(editor);
+      const selectedBlocks = editor.getSelection()?.blocks || [
+        editor.getTextCursorPosition().block,
+      ];
 
-  const fileBlock = useMemo(() => {
-    // Checks if only one block is selected.
-    if (selectedBlocks.length !== 1) {
-      return undefined;
-    }
+      if (selectedBlocks.length !== 1) {
+        return undefined;
+      }
+
+      const block = selectedBlocks[0];
+
+      if (
+        !blockHasType(block, editor, block.type, {
+          url: "string",
+          name: "string",
+        })
+      ) {
+        return undefined;
+      }
+
+      return {
+        blockId: block.id,
+        blockType: block.type,
+        name: block.props.name,
+      };
+    },
+  });
 
-    const block = selectedBlocks[0];
+  const [currentEditingName, setCurrentEditingName] = useState<string>();
 
-    if (
-      blockHasType(block, editor, block.type, {
-        url: "string",
-        name: "string",
-      })
-    ) {
-      setCurrentEditingName(block.props.name);
-      return block;
+  useEffect(() => {
+    if (!state) {
+      return;
     }
+    setCurrentEditingName(state.name);
+  }, [state]);
 
-    return undefined;
-  }, [editor, selectedBlocks]);
+  const handleChange = useCallback(
+    (event: ChangeEvent<HTMLInputElement>) =>
+      setCurrentEditingName(event.currentTarget.value),
+    [],
+  );
 
   const handleEnter = useCallback(
     (event: KeyboardEvent) => {
       if (
-        fileBlock &&
-        editorHasBlockWithType(editor, fileBlock.type, {
+        state !== undefined &&
+        editorHasBlockWithType(editor, state.blockType, {
           name: "string",
         }) &&
         event.key === "Enter"
       ) {
         event.preventDefault();
-        editor.updateBlock(fileBlock, {
+        editor.updateBlock(state.blockId, {
           props: {
             name: currentEditingName,
           },
         });
       }
     },
-    [currentEditingName, editor, fileBlock],
-  );
-
-  const handleChange = useCallback(
-    (event: ChangeEvent<HTMLInputElement>) =>
-      setCurrentEditingName(event.currentTarget.value),
-    [],
+    [currentEditingName, editor, state],
   );
 
-  if (!fileBlock || fileBlock.props.name === "" || !editor.isEditable) {
+  if (state === undefined) {
     return null;
   }
 
@@ -90,11 +108,11 @@ export const FileRenameButton = () => {
         <Components.FormattingToolbar.Button
           className={"bn-button"}
           label={
-            dict.formatting_toolbar.file_rename.tooltip[fileBlock.type] ||
+            dict.formatting_toolbar.file_rename.tooltip[state.blockType] ||
             dict.formatting_toolbar.file_rename.tooltip["file"]
           }
           mainTooltip={
-            dict.formatting_toolbar.file_rename.tooltip[fileBlock.type] ||
+            dict.formatting_toolbar.file_rename.tooltip[state.blockType] ||
             dict.formatting_toolbar.file_rename.tooltip["file"]
           }
           icon={<RiFontFamily />}
@@ -112,7 +130,7 @@ export const FileRenameButton = () => {
             autoFocus={true}
             placeholder={
               dict.formatting_toolbar.file_rename.input_placeholder[
-                fileBlock.type
+                state.blockType
               ] || dict.formatting_toolbar.file_rename.input_placeholder["file"]
             }
             onKeyDown={handleEnter}
diff --git a/packages/react/src/components/FormattingToolbar/DefaultButtons/FileReplaceButton.tsx b/packages/react/src/components/FormattingToolbar/DefaultButtons/FileReplaceButton.tsx
index ed68593186..66e19afaea 100644
--- a/packages/react/src/components/FormattingToolbar/DefaultButtons/FileReplaceButton.tsx
+++ b/packages/react/src/components/FormattingToolbar/DefaultButtons/FileReplaceButton.tsx
@@ -4,12 +4,11 @@ import {
   InlineContentSchema,
   StyleSchema,
 } from "@blocknote/core";
-import { useEffect, useState } from "react";
 import { RiImageEditFill } from "react-icons/ri";
 
 import { useComponentsContext } from "../../../editor/ComponentsContext.js";
 import { useBlockNoteEditor } from "../../../hooks/useBlockNoteEditor.js";
-import { useSelectedBlocks } from "../../../hooks/useSelectedBlocks.js";
+import { useEditorState } from "../../../hooks/useEditorState.js";
 import { useDictionary } from "../../../i18n/dictionary.js";
 import { FilePanel } from "../../FilePanel/FilePanel.js";
 
@@ -23,39 +22,50 @@ export const FileReplaceButton = () => {
     StyleSchema
   >();
 
-  const selectedBlocks = useSelectedBlocks(editor);
+  const state = useEditorState({
+    editor,
+    selector: ({ editor }) => {
+      if (!editor.isEditable) {
+        return undefined;
+      }
 
-  const [isOpen, setIsOpen] = useState<boolean>(false);
+      const selectedBlocks = editor.getSelection()?.blocks || [
+        editor.getTextCursorPosition().block,
+      ];
 
-  useEffect(() => {
-    setIsOpen(false);
-  }, [selectedBlocks]);
+      if (selectedBlocks.length !== 1) {
+        return undefined;
+      }
 
-  const block = selectedBlocks.length === 1 ? selectedBlocks[0] : undefined;
+      const block = selectedBlocks[0];
 
-  if (
-    block === undefined ||
-    !blockHasType(block, editor, block.type, {
-      url: "string",
-    }) ||
-    !editor.isEditable
-  ) {
+      if (
+        !blockHasType(block, editor, block.type, {
+          url: "string",
+        })
+      ) {
+        return undefined;
+      }
+
+      return { blockId: block.id, blockType: block.type };
+    },
+  });
+
+  if (state === undefined) {
     return null;
   }
 
   return (
-    <Components.Generic.Popover.Root opened={isOpen} position={"bottom"}>
+    <Components.Generic.Popover.Root position={"bottom"}>
       <Components.Generic.Popover.Trigger>
         <Components.FormattingToolbar.Button
           className={"bn-button"}
-          onClick={() => setIsOpen(!isOpen)}
-          isSelected={isOpen}
           mainTooltip={
-            dict.formatting_toolbar.file_replace.tooltip[block.type] ||
+            dict.formatting_toolbar.file_replace.tooltip[state.blockType] ||
             dict.formatting_toolbar.file_replace.tooltip["file"]
           }
           label={
-            dict.formatting_toolbar.file_replace.tooltip[block.type] ||
+            dict.formatting_toolbar.file_replace.tooltip[state.blockType] ||
             dict.formatting_toolbar.file_replace.tooltip["file"]
           }
           icon={<RiImageEditFill />}
@@ -65,7 +75,7 @@ export const FileReplaceButton = () => {
         className={"bn-popover-content bn-panel-popover"}
         variant={"panel-popover"}
       >
-        <FilePanel block={block as any} />
+        <FilePanel blockId={state.blockId} />
       </Components.Generic.Popover.Content>
     </Components.Generic.Popover.Root>
   );
diff --git a/packages/react/src/components/FormattingToolbar/DefaultButtons/NestBlockButtons.tsx b/packages/react/src/components/FormattingToolbar/DefaultButtons/NestBlockButtons.tsx
index 9e42657899..7bc87c45ec 100644
--- a/packages/react/src/components/FormattingToolbar/DefaultButtons/NestBlockButtons.tsx
+++ b/packages/react/src/components/FormattingToolbar/DefaultButtons/NestBlockButtons.tsx
@@ -4,13 +4,12 @@ import {
   InlineContentSchema,
   StyleSchema,
 } from "@blocknote/core";
-import { useCallback, useMemo, useState } from "react";
+import { useCallback } from "react";
 import { RiIndentDecrease, RiIndentIncrease } from "react-icons/ri";
 
 import { useComponentsContext } from "../../../editor/ComponentsContext.js";
 import { useBlockNoteEditor } from "../../../hooks/useBlockNoteEditor.js";
-import { useEditorContentOrSelectionChange } from "../../../hooks/useEditorContentOrSelectionChange.js";
-import { useSelectedBlocks } from "../../../hooks/useSelectedBlocks.js";
+import { useEditorState } from "../../../hooks/useEditorState.js";
 import { useDictionary } from "../../../i18n/dictionary.js";
 
 export const NestBlockButton = () => {
@@ -23,28 +22,37 @@ export const NestBlockButton = () => {
     StyleSchema
   >();
 
-  const selectedBlocks = useSelectedBlocks(editor);
-
-  const [canNestBlock, setCanNestBlock] = useState<boolean>(() =>
-    editor.canNestBlock(),
-  );
-
-  useEditorContentOrSelectionChange(() => {
-    setCanNestBlock(editor.canNestBlock());
-  }, editor);
+  const state = useEditorState({
+    editor,
+    selector: ({ editor }) => {
+      // Do not show if:
+      if (
+        // The editor is read-only.
+        !editor.isEditable ||
+        // None of the selected blocks have inline content
+        !(
+          editor.getSelection()?.blocks || [
+            editor.getTextCursorPosition().block,
+          ]
+        ).find((block) => block.content !== undefined)
+      ) {
+        return undefined;
+      }
+
+      return {
+        canNestBlock: editor.canNestBlock(),
+      };
+    },
+  });
 
   const nestBlock = useCallback(() => {
-    editor.focus();
-    editor.nestBlock();
-  }, [editor]);
-
-  const show = useMemo(() => {
-    return !selectedBlocks.find(
-      (block) => editor.schema.blockSchema[block.type].content !== "inline",
-    );
-  }, [editor.schema.blockSchema, selectedBlocks]);
+    if (state !== undefined && state.canNestBlock) {
+      editor.focus();
+      editor.nestBlock();
+    }
+  }, [editor, state]);
 
-  if (!show || !editor.isEditable) {
+  if (state === undefined) {
     return null;
   }
 
@@ -53,7 +61,7 @@ export const NestBlockButton = () => {
       className={"bn-button"}
       data-test="nestBlock"
       onClick={nestBlock}
-      isDisabled={!canNestBlock}
+      isDisabled={!state.canNestBlock}
       label={dict.formatting_toolbar.nest.tooltip}
       mainTooltip={dict.formatting_toolbar.nest.tooltip}
       secondaryTooltip={formatKeyboardShortcut(
@@ -71,28 +79,37 @@ export const UnnestBlockButton = () => {
 
   const editor = useBlockNoteEditor<any, any, any>();
 
-  const selectedBlocks = useSelectedBlocks(editor);
-
-  const [canUnnestBlock, setCanUnnestBlock] = useState<boolean>(() =>
-    editor.canUnnestBlock(),
-  );
-
-  useEditorContentOrSelectionChange(() => {
-    setCanUnnestBlock(editor.canUnnestBlock());
-  }, editor);
+  const state = useEditorState({
+    editor,
+    selector: ({ editor }) => {
+      // Do not show if:
+      if (
+        // The editor is read-only.
+        !editor.isEditable ||
+        // None of the selected blocks have inline content
+        !(
+          editor.getSelection()?.blocks || [
+            editor.getTextCursorPosition().block,
+          ]
+        ).find((block) => block.content !== undefined)
+      ) {
+        return undefined;
+      }
+
+      return {
+        canUnnestBlock: editor.canUnnestBlock(),
+      };
+    },
+  });
 
   const unnestBlock = useCallback(() => {
-    editor.focus();
-    editor.unnestBlock();
-  }, [editor]);
-
-  const show = useMemo(() => {
-    return !selectedBlocks.find(
-      (block) => editor.schema.blockSchema[block.type].content !== "inline",
-    );
-  }, [editor.schema.blockSchema, selectedBlocks]);
+    if (state !== undefined && state.canUnnestBlock) {
+      editor.focus();
+      editor.unnestBlock();
+    }
+  }, [editor, state]);
 
-  if (!show || !editor.isEditable) {
+  if (state === undefined) {
     return null;
   }
 
@@ -101,7 +118,7 @@ export const UnnestBlockButton = () => {
       className={"bn-button"}
       data-test="unnestBlock"
       onClick={unnestBlock}
-      isDisabled={!canUnnestBlock}
+      isDisabled={!state.canUnnestBlock}
       label={dict.formatting_toolbar.unnest.tooltip}
       mainTooltip={dict.formatting_toolbar.unnest.tooltip}
       secondaryTooltip={formatKeyboardShortcut(
diff --git a/packages/react/src/components/FormattingToolbar/DefaultButtons/TableCellMergeButton.tsx b/packages/react/src/components/FormattingToolbar/DefaultButtons/TableCellMergeButton.tsx
index bb11326a08..81567c5522 100644
--- a/packages/react/src/components/FormattingToolbar/DefaultButtons/TableCellMergeButton.tsx
+++ b/packages/react/src/components/FormattingToolbar/DefaultButtons/TableCellMergeButton.tsx
@@ -3,13 +3,15 @@ import {
   InlineContentSchema,
   StyleSchema,
 } from "@blocknote/core";
-import { useCallback, useMemo } from "react";
+import { TableHandles } from "@blocknote/core/extensions";
+import { useCallback } from "react";
 import { RiMergeCellsHorizontal, RiMergeCellsVertical } from "react-icons/ri";
 
 import { useComponentsContext } from "../../../editor/ComponentsContext.js";
 import { useBlockNoteEditor } from "../../../hooks/useBlockNoteEditor.js";
-import { useSelectedBlocks } from "../../../hooks/useSelectedBlocks.js";
+import { useEditorState } from "../../../hooks/useEditorState.js";
 import { useDictionary } from "../../../i18n/dictionary.js";
+import { useExtension } from "../../../hooks/useExtension.js";
 
 export const TableCellMergeButton = () => {
   const dict = useDictionary();
@@ -23,31 +25,39 @@ export const TableCellMergeButton = () => {
     StyleSchema
   >();
 
-  const selectedBlocks = useSelectedBlocks(editor);
-  const mergeDirection = useMemo(() => {
-    // Checks if only one block is selected.
-    if (selectedBlocks.length !== 1) {
-      return undefined;
-    }
+  const tableHandles = useExtension(TableHandles);
+  const state = useEditorState({
+    editor,
+    selector: ({ editor }) => {
+      if (!editor.isEditable || !editor.settings.tables.splitCells) {
+        return undefined;
+      }
+
+      const selectedBlocks = editor.getSelection()?.blocks || [
+        editor.getTextCursorPosition().block,
+      ];
+
+      if (selectedBlocks.length !== 1) {
+        return undefined;
+      }
 
-    const block = selectedBlocks[0];
+      const block = selectedBlocks[0];
 
-    if (block.type === "table") {
-      return editor.tableHandles?.getMergeDirection(block);
-    }
+      if (block.type !== "table") {
+        return undefined;
+      }
 
-    return undefined;
-  }, [editor, selectedBlocks]);
+      return {
+        mergeDirection: tableHandles.getMergeDirection(block),
+      };
+    },
+  });
 
   const onClick = useCallback(() => {
-    editor.tableHandles?.mergeCells();
-  }, [editor]);
+    tableHandles?.mergeCells();
+  }, [tableHandles]);
 
-  if (
-    !editor.isEditable ||
-    mergeDirection === undefined ||
-    !editor.settings.tables.splitCells
-  ) {
+  if (state === undefined) {
     return null;
   }
 
@@ -57,7 +67,7 @@ export const TableCellMergeButton = () => {
       label={dict.formatting_toolbar.table_cell_merge.tooltip}
       mainTooltip={dict.formatting_toolbar.table_cell_merge.tooltip}
       icon={
-        mergeDirection === "horizontal" ? (
+        state.mergeDirection === "horizontal" ? (
           <RiMergeCellsHorizontal />
         ) : (
           <RiMergeCellsVertical />
diff --git a/packages/react/src/components/FormattingToolbar/DefaultButtons/TextAlignButton.tsx b/packages/react/src/components/FormattingToolbar/DefaultButtons/TextAlignButton.tsx
index 05b38f6a7a..fff852769e 100644
--- a/packages/react/src/components/FormattingToolbar/DefaultButtons/TextAlignButton.tsx
+++ b/packages/react/src/components/FormattingToolbar/DefaultButtons/TextAlignButton.tsx
@@ -9,7 +9,8 @@ import {
   StyleSchema,
   TableContent,
 } from "@blocknote/core";
-import { useCallback, useMemo } from "react";
+import { TableHandles } from "@blocknote/core/extensions";
+import { useCallback } from "react";
 import { IconType } from "react-icons";
 import {
   RiAlignCenter,
@@ -20,7 +21,8 @@ import {
 
 import { useComponentsContext } from "../../../editor/ComponentsContext.js";
 import { useBlockNoteEditor } from "../../../hooks/useBlockNoteEditor.js";
-import { useSelectedBlocks } from "../../../hooks/useSelectedBlocks.js";
+import { useEditorState } from "../../../hooks/useEditorState.js";
+import { useExtension } from "../../../hooks/useExtension.js";
 import { useDictionary } from "../../../i18n/dictionary.js";
 
 type TextAlignment = DefaultProps["textAlignment"];
@@ -42,44 +44,62 @@ export const TextAlignButton = (props: { textAlignment: TextAlignment }) => {
     StyleSchema
   >();
 
-  const selectedBlocks = useSelectedBlocks(editor);
-
-  const textAlignment = useMemo(() => {
-    const block = selectedBlocks[0];
-
-    if (
-      blockHasType(block, editor, block.type, {
-        textAlignment: defaultProps.textAlignment,
-      })
-    ) {
-      return block.props.textAlignment;
-    }
-    if (block.type === "table") {
-      const cellSelection = editor.tableHandles?.getCellSelection();
-      if (!cellSelection) {
-        return;
+  const tableHandles = useExtension(TableHandles);
+
+  const state = useEditorState({
+    editor,
+    selector: ({ editor }) => {
+      if (!editor.isEditable) {
+        return undefined;
       }
-      const allCellsInTable = cellSelection.cells.map(
-        ({ row, col }) =>
-          mapTableCell(
-            (block.content as TableContent<any, any>).rows[row].cells[col],
-          ).props.textAlignment,
-      );
-      const firstAlignment = allCellsInTable[0];
 
-      if (allCellsInTable.every((alignment) => alignment === firstAlignment)) {
-        return firstAlignment;
+      const selectedBlocks = editor.getSelection()?.blocks || [
+        editor.getTextCursorPosition().block,
+      ];
+
+      const firstBlock = selectedBlocks[0];
+
+      if (
+        blockHasType(firstBlock, editor, firstBlock.type, {
+          textAlignment: defaultProps.textAlignment,
+        })
+      ) {
+        return {
+          textAlignment: firstBlock.props.textAlignment,
+          blocks: selectedBlocks,
+        };
       }
-    }
 
-    return;
-  }, [editor, selectedBlocks]);
+      if (
+        selectedBlocks.length === 1 &&
+        blockHasType(firstBlock, editor, "table")
+      ) {
+        const cellSelection = tableHandles.getCellSelection();
+        if (!cellSelection) {
+          return undefined;
+        }
+
+        return {
+          textAlignment: mapTableCell(
+            (firstBlock.content as TableContent<any, any>).rows[0].cells[0],
+          ).props.textAlignment,
+          blocks: [firstBlock],
+        };
+      }
+
+      return undefined;
+    },
+  });
 
   const setTextAlignment = useCallback(
     (textAlignment: TextAlignment) => {
+      if (state === undefined) {
+        return;
+      }
+
       editor.focus();
 
-      for (const block of selectedBlocks) {
+      for (const block of state.blocks) {
         if (
           blockHasType(block, editor, block.type, {
             textAlignment: defaultProps.textAlignment,
@@ -92,7 +112,7 @@ export const TextAlignButton = (props: { textAlignment: TextAlignment }) => {
             props: { textAlignment: textAlignment },
           });
         } else if (block.type === "table") {
-          const cellSelection = editor.tableHandles?.getCellSelection();
+          const cellSelection = tableHandles.getCellSelection();
           if (!cellSelection) {
             continue;
           }
@@ -128,20 +148,10 @@ export const TextAlignButton = (props: { textAlignment: TextAlignment }) => {
         }
       }
     },
-    [editor, selectedBlocks],
+    [editor, state, tableHandles],
   );
 
-  const show = useMemo(() => {
-    return !!selectedBlocks.find(
-      (block) =>
-        blockHasType(block, editor, block.type, {
-          textAlignment: defaultProps.textAlignment,
-        }) ||
-        (block.type === "table" && block.children),
-    );
-  }, [editor, selectedBlocks]);
-
-  if (!show || !editor.isEditable) {
+  if (state === undefined) {
     return null;
   }
 
@@ -154,7 +164,7 @@ export const TextAlignButton = (props: { textAlignment: TextAlignment }) => {
         props.textAlignment.slice(1)
       }`}
       onClick={() => setTextAlignment(props.textAlignment)}
-      isSelected={textAlignment === props.textAlignment}
+      isSelected={state.textAlignment === props.textAlignment}
       label={dict.formatting_toolbar[`align_${props.textAlignment}`].tooltip}
       mainTooltip={
         dict.formatting_toolbar[`align_${props.textAlignment}`].tooltip
diff --git a/packages/react/src/components/FormattingToolbar/DefaultSelects/BlockTypeSelect.tsx b/packages/react/src/components/FormattingToolbar/DefaultSelects/BlockTypeSelect.tsx
index 3aa6f3165d..d9bb12a00a 100644
--- a/packages/react/src/components/FormattingToolbar/DefaultSelects/BlockTypeSelect.tsx
+++ b/packages/react/src/components/FormattingToolbar/DefaultSelects/BlockTypeSelect.tsx
@@ -27,7 +27,7 @@ import {
   useComponentsContext,
 } from "../../../editor/ComponentsContext.js";
 import { useBlockNoteEditor } from "../../../hooks/useBlockNoteEditor.js";
-import { useSelectedBlocks } from "../../../hooks/useSelectedBlocks.js";
+import { useEditorState } from "../../../hooks/useEditorState.js";
 
 export type BlockTypeSelectItem = {
   name: string;
@@ -134,7 +134,11 @@ export const BlockTypeSelect = (props: { items?: BlockTypeSelectItem[] }) => {
     StyleSchema
   >();
 
-  const selectedBlocks = useSelectedBlocks(editor);
+  const selectedBlocks = useEditorState({
+    editor,
+    selector: ({ editor }) =>
+      editor.getSelection()?.blocks || [editor.getTextCursorPosition().block],
+  });
   const firstSelectedBlock = selectedBlocks[0];
 
   // Filters out all items in which the block type and props don't conform to
diff --git a/packages/react/src/components/FormattingToolbar/ExperimentalMobileFormattingToolbarController.tsx b/packages/react/src/components/FormattingToolbar/ExperimentalMobileFormattingToolbarController.tsx
index ed997cf33a..6a2caa80c5 100644
--- a/packages/react/src/components/FormattingToolbar/ExperimentalMobileFormattingToolbarController.tsx
+++ b/packages/react/src/components/FormattingToolbar/ExperimentalMobileFormattingToolbarController.tsx
@@ -1,8 +1,9 @@
 import { BlockSchema, InlineContentSchema, StyleSchema } from "@blocknote/core";
-import { UseFloatingOptions } from "@floating-ui/react";
+import { FormattingToolbar as FormattingToolbarExtension } from "@blocknote/core/extensions";
 import { FC, CSSProperties, useMemo, useRef, useState, useEffect } from "react";
+
 import { useBlockNoteEditor } from "../../hooks/useBlockNoteEditor.js";
-import { useUIPluginState } from "../../hooks/useUIPluginState.js";
+import { useExtensionState } from "../../hooks/useExtension.js";
 import { FormattingToolbar } from "./FormattingToolbar.js";
 import { FormattingToolbarProps } from "./FormattingToolbarProps.js";
 
@@ -14,7 +15,6 @@ import { FormattingToolbarProps } from "./FormattingToolbarProps.js";
  */
 export const ExperimentalMobileFormattingToolbarController = (props: {
   formattingToolbar?: FC<FormattingToolbarProps>;
-  floatingOptions?: Partial<UseFloatingOptions>;
 }) => {
   const [transform, setTransform] = useState<string>("none");
   const divRef = useRef<HTMLDivElement>(null);
@@ -23,15 +23,18 @@ export const ExperimentalMobileFormattingToolbarController = (props: {
     InlineContentSchema,
     StyleSchema
   >();
-  const state = useUIPluginState(
-    editor.formattingToolbar.onUpdate.bind(editor.formattingToolbar),
-  );
+
+  const show = useExtensionState(FormattingToolbarExtension, {
+    editor,
+    selector: (state) => state,
+  });
+
   const style = useMemo<CSSProperties>(() => {
     return {
       display: "flex",
       position: "fixed",
       bottom: 0,
-      zIndex: 3000,
+      zIndex: `calc(var(--bn-ui-base-z-index) + 40)`,
       transform,
     };
   }, [transform]);
@@ -63,11 +66,7 @@ export const ExperimentalMobileFormattingToolbarController = (props: {
     };
   }, []);
 
-  if (!state) {
-    return null;
-  }
-
-  if (!state.show && divRef.current) {
+  if (!show && divRef.current) {
     // The component is fading out. Use the previous state to render the toolbar with innerHTML,
     // because otherwise the toolbar will quickly flickr (i.e.: show a different state) while fading out,
     // which looks weird
diff --git a/packages/react/src/components/FormattingToolbar/FormattingToolbarController.tsx b/packages/react/src/components/FormattingToolbar/FormattingToolbarController.tsx
index 6ce960d035..b2950daff5 100644
--- a/packages/react/src/components/FormattingToolbar/FormattingToolbarController.tsx
+++ b/packages/react/src/components/FormattingToolbar/FormattingToolbarController.tsx
@@ -1,18 +1,20 @@
 import {
+  blockHasType,
   BlockSchema,
+  defaultProps,
   DefaultProps,
   InlineContentSchema,
   StyleSchema,
 } from "@blocknote/core";
-import { UseFloatingOptions, flip, offset, shift } from "@floating-ui/react";
-import { isEventTargetWithin } from "@floating-ui/react/utils";
-import { FC, useMemo, useRef, useState } from "react";
+import { FormattingToolbar as FormattingToolbarExtension } from "@blocknote/core/extensions";
+import { flip, offset, shift } from "@floating-ui/react";
+import { FC, useMemo } from "react";
 
 import { useBlockNoteEditor } from "../../hooks/useBlockNoteEditor.js";
-import { useEditorContentOrSelectionChange } from "../../hooks/useEditorContentOrSelectionChange.js";
-import { useUIElementPositioning } from "../../hooks/useUIElementPositioning.js";
-import { useUIPluginState } from "../../hooks/useUIPluginState.js";
-import { mergeRefs } from "../../util/mergeRefs.js";
+import { useEditorState } from "../../hooks/useEditorState.js";
+import { useExtension, useExtensionState } from "../../hooks/useExtension.js";
+import { FloatingUIOptions } from "../Popovers/FloatingUIOptions.js";
+import { PositionPopover } from "../Popovers/PositionPopover.js";
 import { FormattingToolbar } from "./FormattingToolbar.js";
 import { FormattingToolbarProps } from "./FormattingToolbarProps.js";
 
@@ -33,107 +35,78 @@ const textAlignmentToPlacement = (
 
 export const FormattingToolbarController = (props: {
   formattingToolbar?: FC<FormattingToolbarProps>;
-  floatingOptions?: Partial<UseFloatingOptions>;
+  floatingUIOptions?: FloatingUIOptions;
 }) => {
-  const divRef = useRef<HTMLDivElement>(null);
-
   const editor = useBlockNoteEditor<
     BlockSchema,
     InlineContentSchema,
     StyleSchema
   >();
+  const formattingToolbar = useExtension(FormattingToolbarExtension, {
+    editor,
+  });
+  const show = useExtensionState(FormattingToolbarExtension, {
+    editor,
+  });
+
+  const position = useEditorState({
+    editor,
+    selector: ({ editor }) =>
+      formattingToolbar.store.state
+        ? {
+            from: editor.prosemirrorState.selection.from,
+            to: editor.prosemirrorState.selection.to,
+          }
+        : undefined,
+  });
 
-  const [placement, setPlacement] = useState<"top-start" | "top" | "top-end">(
-    () => {
+  const placement = useEditorState({
+    editor,
+    selector: ({ editor }) => {
       const block = editor.getTextCursorPosition().block;
 
-      if (!("textAlignment" in block.props)) {
+      if (
+        !blockHasType(block, editor, block.type, {
+          textAlignment: defaultProps.textAlignment,
+        })
+      ) {
         return "top-start";
+      } else {
+        return textAlignmentToPlacement(block.props.textAlignment);
       }
-
-      return textAlignmentToPlacement(
-        block.props.textAlignment as DefaultProps["textAlignment"],
-      );
     },
-  );
-
-  useEditorContentOrSelectionChange(() => {
-    const block = editor.getTextCursorPosition().block;
-
-    if (!("textAlignment" in block.props)) {
-      setPlacement("top-start");
-    } else {
-      setPlacement(
-        textAlignmentToPlacement(
-          block.props.textAlignment as DefaultProps["textAlignment"],
-        ),
-      );
-    }
-  }, editor);
-
-  const state = useUIPluginState(
-    editor.formattingToolbar.onUpdate.bind(editor.formattingToolbar),
-  );
-
-  const { isMounted, ref, style, getFloatingProps } = useUIElementPositioning(
-    state?.show || false,
-    state?.referencePos || null,
-    3000,
-    {
-      placement,
-      middleware: [offset(10), shift(), flip()],
-      onOpenChange: (open, _event) => {
-        // console.log("change", event);
-        if (!open) {
-          editor.formattingToolbar.closeMenu();
-          editor.focus();
-        }
+  });
+
+  const floatingUIOptions = useMemo<FloatingUIOptions>(
+    () => ({
+      useFloatingOptions: {
+        open: show,
+        // Needed as hooks like `useDismiss` call `onOpenChange` to change the
+        // open state.
+        onOpenChange: (open) => {
+          formattingToolbar.store.setState(open);
+        },
+        placement,
+        middleware: [offset(10), shift(), flip()],
       },
-      canDismiss: {
-        enabled: true,
-        escapeKey: true,
-        outsidePress: (e) => {
-          const view = editor._tiptapEditor?.view;
-          if (!view) {
-            return false;
-          }
-
-          const target = e.target;
-          if (!target) {
-            return false;
-          }
-
-          return !isEventTargetWithin(e, view.dom.parentElement);
+      elementProps: {
+        style: {
+          zIndex: 40,
         },
       },
-      ...props.floatingOptions,
-    },
+    }),
+    [formattingToolbar.store, show, placement],
   );
 
-  const combinedRef = useMemo(() => mergeRefs([divRef, ref]), [divRef, ref]);
-
-  if (!isMounted || !state) {
-    return null;
-  }
-
-  if (!state.show && divRef.current) {
-    // The component is fading out. Use the previous state to render the toolbar with innerHTML,
-    // because otherwise the toolbar will quickly flickr (i.e.: show a different state) while fading out,
-    // which looks weird
-    return (
-      <div
-        ref={combinedRef}
-        style={style}
-        dangerouslySetInnerHTML={{ __html: divRef.current.innerHTML }}
-      ></div>
-    );
-  }
-
   const Component = props.formattingToolbar || FormattingToolbar;
 
   return (
-    <div ref={combinedRef} style={style} {...getFloatingProps()}>
-      <Component />
-    </div>
+    <PositionPopover
+      position={position}
+      {...floatingUIOptions}
+      {...props.floatingUIOptions}
+    >
+      {show && <Component />}
+    </PositionPopover>
   );
 };
diff --git a/packages/react/src/components/LinkToolbar/DefaultButtons/DeleteLinkButton.tsx b/packages/react/src/components/LinkToolbar/DefaultButtons/DeleteLinkButton.tsx
index b6fe8f53aa..40981752f8 100644
--- a/packages/react/src/components/LinkToolbar/DefaultButtons/DeleteLinkButton.tsx
+++ b/packages/react/src/components/LinkToolbar/DefaultButtons/DeleteLinkButton.tsx
@@ -2,20 +2,22 @@ import { RiLinkUnlink } from "react-icons/ri";
 
 import { useComponentsContext } from "../../../editor/ComponentsContext.js";
 import { useDictionary } from "../../../i18n/dictionary.js";
-import { LinkToolbarProps } from "../LinkToolbarProps.js";
+import { useExtension } from "../../../hooks/useExtension.js";
+import { LinkToolbar } from "@blocknote/core/extensions";
 
-export const DeleteLinkButton = (
-  props: Pick<LinkToolbarProps, "deleteLink">,
-) => {
+export const DeleteLinkButton = () => {
   const Components = useComponentsContext()!;
   const dict = useDictionary();
+
+  const { deleteLink } = useExtension(LinkToolbar);
+
   return (
     <Components.LinkToolbar.Button
       className={"bn-button"}
       label={dict.link_toolbar.delete.tooltip}
       mainTooltip={dict.link_toolbar.delete.tooltip}
       isSelected={false}
-      onClick={props.deleteLink}
+      onClick={deleteLink}
       icon={<RiLinkUnlink />}
     />
   );
diff --git a/packages/react/src/components/LinkToolbar/DefaultButtons/EditLinkButton.tsx b/packages/react/src/components/LinkToolbar/DefaultButtons/EditLinkButton.tsx
index f83ecb5adb..147cf4f526 100644
--- a/packages/react/src/components/LinkToolbar/DefaultButtons/EditLinkButton.tsx
+++ b/packages/react/src/components/LinkToolbar/DefaultButtons/EditLinkButton.tsx
@@ -4,7 +4,7 @@ import { EditLinkMenuItems } from "../EditLinkMenuItems.js";
 import { LinkToolbarProps } from "../LinkToolbarProps.js";
 
 export const EditLinkButton = (
-  props: Pick<LinkToolbarProps, "url" | "text" | "editLink">,
+  props: Pick<LinkToolbarProps, "url" | "text" | "range">,
 ) => {
   const Components = useComponentsContext()!;
   const dict = useDictionary();
diff --git a/packages/react/src/components/LinkToolbar/EditLinkMenuItems.tsx b/packages/react/src/components/LinkToolbar/EditLinkMenuItems.tsx
index 9f59c1272d..40b6b212b4 100644
--- a/packages/react/src/components/LinkToolbar/EditLinkMenuItems.tsx
+++ b/packages/react/src/components/LinkToolbar/EditLinkMenuItems.tsx
@@ -1,4 +1,8 @@
-import { DEFAULT_LINK_PROTOCOL, VALID_LINK_PROTOCOLS } from "@blocknote/core";
+import {
+  DEFAULT_LINK_PROTOCOL,
+  LinkToolbar,
+  VALID_LINK_PROTOCOLS,
+} from "@blocknote/core/extensions";
 import {
   ChangeEvent,
   KeyboardEvent,
@@ -8,6 +12,7 @@ import {
 } from "react";
 import { RiLink, RiText } from "react-icons/ri";
 import { useComponentsContext } from "../../editor/ComponentsContext.js";
+import { useExtension } from "../../hooks/useExtension.js";
 import { useDictionary } from "../../i18n/dictionary.js";
 import { LinkToolbarProps } from "./LinkToolbarProps.js";
 
@@ -22,14 +27,16 @@ const validateUrl = (url: string) => {
 };
 
 export const EditLinkMenuItems = (
-  props: Pick<LinkToolbarProps, "url" | "text" | "editLink"> & {
+  props: Pick<LinkToolbarProps, "url" | "text" | "range"> & {
     showTextField?: boolean;
   },
 ) => {
   const Components = useComponentsContext()!;
   const dict = useDictionary();
 
-  const { url, text, editLink, showTextField } = props;
+  const { editLink } = useExtension(LinkToolbar);
+
+  const { url, text, showTextField } = props;
 
   const [currentUrl, setCurrentUrl] = useState<string>(url);
   const [currentText, setCurrentText] = useState<string>(text);
@@ -43,10 +50,10 @@ export const EditLinkMenuItems = (
     (event: KeyboardEvent) => {
       if (event.key === "Enter") {
         event.preventDefault();
-        editLink(validateUrl(currentUrl), currentText);
+        editLink(validateUrl(currentUrl), currentText, props.range.from);
       }
     },
-    [editLink, currentUrl, currentText],
+    [editLink, currentUrl, currentText, props.range.from],
   );
 
   const handleUrlChange = useCallback(
@@ -62,8 +69,8 @@ export const EditLinkMenuItems = (
   );
 
   const handleSubmit = useCallback(
-    () => editLink(validateUrl(currentUrl), currentText),
-    [editLink, currentUrl, currentText],
+    () => editLink(validateUrl(currentUrl), currentText, props.range.from),
+    [editLink, currentUrl, currentText, props.range],
   );
 
   return (
diff --git a/packages/react/src/components/LinkToolbar/LinkToolbar.tsx b/packages/react/src/components/LinkToolbar/LinkToolbar.tsx
index 8f66c68a55..af179b9cca 100644
--- a/packages/react/src/components/LinkToolbar/LinkToolbar.tsx
+++ b/packages/react/src/components/LinkToolbar/LinkToolbar.tsx
@@ -1,5 +1,3 @@
-import { ReactNode } from "react";
-
 import { useComponentsContext } from "../../editor/ComponentsContext.js";
 import { DeleteLinkButton } from "./DefaultButtons/DeleteLinkButton.js";
 import { EditLinkButton } from "./DefaultButtons/EditLinkButton.js";
@@ -17,26 +15,20 @@ import { LinkToolbarProps } from "./LinkToolbarProps.js";
  * - Custom buttons: The `ToolbarButton` component in the
  * `components/mantine-shared/Toolbar` directory.
  */
-export const LinkToolbar = (
-  props: LinkToolbarProps & { children?: ReactNode },
-) => {
+export const LinkToolbar = (props: LinkToolbarProps) => {
   const Components = useComponentsContext()!;
 
   return (
-    <Components.LinkToolbar.Root
-      className={"bn-toolbar bn-link-toolbar"}
-      onMouseEnter={props.stopHideTimer}
-      onMouseLeave={props.startHideTimer}
-    >
+    <Components.LinkToolbar.Root className={"bn-toolbar bn-link-toolbar"}>
       {props.children || (
         <>
           <EditLinkButton
             url={props.url}
             text={props.text}
-            editLink={props.editLink}
+            range={props.range}
           />
           <OpenLinkButton url={props.url} />
-          <DeleteLinkButton deleteLink={props.deleteLink} />
+          <DeleteLinkButton />
         </>
       )}
     </Components.LinkToolbar.Root>
diff --git a/packages/react/src/components/LinkToolbar/LinkToolbarController.tsx b/packages/react/src/components/LinkToolbar/LinkToolbarController.tsx
index dc0fbf4a5e..a3a93dd9ca 100644
--- a/packages/react/src/components/LinkToolbar/LinkToolbarController.tsx
+++ b/packages/react/src/components/LinkToolbar/LinkToolbarController.tsx
@@ -1,68 +1,143 @@
-import {
-  BlockSchema,
-  DefaultBlockSchema,
-  DefaultInlineContentSchema,
-  DefaultStyleSchema,
-  InlineContentSchema,
-  StyleSchema,
-} from "@blocknote/core";
-import { UseFloatingOptions, flip, offset } from "@floating-ui/react";
-import { FC } from "react";
+import { BlockSchema, InlineContentSchema, StyleSchema } from "@blocknote/core";
+import { LinkToolbar as LinkToolbarExtension } from "@blocknote/core/extensions";
+import { flip, offset, safePolygon } from "@floating-ui/react";
+import { Range } from "@tiptap/core";
+import { FC, useEffect, useMemo, useState } from "react";
 
 import { useBlockNoteEditor } from "../../hooks/useBlockNoteEditor.js";
-import { useUIElementPositioning } from "../../hooks/useUIElementPositioning.js";
-import { useUIPluginState } from "../../hooks/useUIPluginState.js";
 import { LinkToolbar } from "./LinkToolbar.js";
 import { LinkToolbarProps } from "./LinkToolbarProps.js";
+import { useEditorState } from "../../hooks/useEditorState.js";
+import { useExtension } from "../../hooks/useExtension.js";
+import { GenericPopover } from "../Popovers/GenericPopover.js";
+import { FloatingUIOptions } from "../Popovers/FloatingUIOptions.js";
 
-export const LinkToolbarController = <
-  BSchema extends BlockSchema = DefaultBlockSchema,
-  I extends InlineContentSchema = DefaultInlineContentSchema,
-  S extends StyleSchema = DefaultStyleSchema,
->(props: {
+export const LinkToolbarController = (props: {
   linkToolbar?: FC<LinkToolbarProps>;
-  floatingOptions?: Partial<UseFloatingOptions>;
+  floatingUIOptions?: FloatingUIOptions;
 }) => {
-  const editor = useBlockNoteEditor<BSchema, I, S>();
+  const editor = useBlockNoteEditor<
+    BlockSchema,
+    InlineContentSchema,
+    StyleSchema
+  >();
 
-  const callbacks = {
-    deleteLink: editor.linkToolbar.deleteLink,
-    editLink: editor.linkToolbar.editLink,
-    startHideTimer: editor.linkToolbar.startHideTimer,
-    stopHideTimer: editor.linkToolbar.stopHideTimer,
-  };
+  const [open, setOpen] = useState(false);
 
-  const state = useUIPluginState(
-    editor.linkToolbar.onUpdate.bind(editor.linkToolbar),
-  );
-  const { isMounted, ref, style, getFloatingProps } = useUIElementPositioning(
-    state?.show || false,
-    state?.referencePos || null,
-    4000,
-    {
-      placement: "top-start",
-      middleware: [offset(10), flip()],
-      onOpenChange: (open) => {
-        if (!open) {
-          editor.linkToolbar.closeMenu();
-          editor.focus();
-        }
-      },
-      ...props.floatingOptions,
+  const linkToolbar = useExtension(LinkToolbarExtension);
+  const selectionLink = useEditorState({
+    editor,
+    selector: () => {
+      const link = linkToolbar.getLinkAtSelection();
+      if (!link) {
+        return undefined;
+      }
+
+      return {
+        url: link.mark.attrs.href as string,
+        text: link.text,
+        range: link.range,
+        element: linkToolbar.getLinkElementAtPos(link.range.from)!,
+      };
     },
+  });
+  useEffect(() => {
+    setOpen(!!selectionLink);
+    if (selectionLink) {
+      // Clears the link hovered by the mouse cursor, when the text cursor is
+      // within a link, to avoid any potential clashes in positioning.
+      setMouseHoverLink(undefined);
+    }
+  }, [selectionLink]);
+
+  // The `mouseHoverLink` state is completely separate from the `open` state as
+  // the FloatingUI `useHover` hook handles opening/closing the popover when a
+  // link is hovered with the mouse cursor. Therefore, we only need to update
+  // the link when a new one is hovered.
+  const [mouseHoverLink, setMouseHoverLink] = useState<
+    | { url: string; text: string; range: Range; element: HTMLAnchorElement }
+    | undefined
+  >(undefined);
+  useEffect(() => {
+    const cb = (event: MouseEvent) => {
+      // Ignores the link hovered by the mouse cursor, when the text cursor is
+      // within a link, to avoid any potential clashes in positioning.
+      if (selectionLink) {
+        return;
+      }
+
+      if (!(event.target instanceof HTMLElement)) {
+        return;
+      }
+
+      const link = linkToolbar.getLinkAtElement(event.target);
+      if (!link) {
+        return;
+      }
+
+      setMouseHoverLink({
+        url: link.mark.attrs.href as string,
+        text: link.text,
+        range: link.range,
+        element: linkToolbar.getLinkElementAtPos(link.range.from)!,
+      });
+    };
+
+    document.body.addEventListener("mouseover", cb);
+
+    return () => {
+      document.body.removeEventListener("mouseover", cb);
+    };
+  }, [linkToolbar, mouseHoverLink?.url, selectionLink]);
+
+  const link = useMemo(
+    () => selectionLink || mouseHoverLink,
+    [mouseHoverLink, selectionLink],
   );
 
-  if (!isMounted || !state) {
-    return null;
-  }
+  const floatingUIOptions = useMemo<FloatingUIOptions>(
+    () => ({
+      useFloatingOptions: {
+        open,
+        onOpenChange: (open, _event, reason) => {
+          // We want to prioritize `selectionLink` over `mouseHoverLink`, so we
+          // ignore opening/closing from hover events.
+          if (selectionLink && reason === "hover") {
+            return;
+          }
 
-  const { show, referencePos, ...data } = state;
+          setOpen(open);
+        },
+        placement: "top-start",
+        middleware: [offset(10), flip()],
+      },
+      useHoverProps: {
+        // `useHover` hook only enabled when a link is hovered with the
+        // mouse.
+        enabled: !selectionLink && !!mouseHoverLink,
+        delay: {
+          // TODO: Figure out good values
+          open: 250,
+          close: 250,
+        },
+        // restMs: 300,
+        handleClose: safePolygon(),
+      },
+      elementProps: {
+        style: {
+          zIndex: 50,
+        },
+      },
+      ...props.floatingUIOptions,
+    }),
+    [mouseHoverLink, open, props.floatingUIOptions, selectionLink],
+  );
 
   const Component = props.linkToolbar || LinkToolbar;
 
   return (
-    <div ref={ref} style={style} {...getFloatingProps()}>
-      <Component {...data} {...callbacks} />
-    </div>
+    <GenericPopover reference={link?.element} {...floatingUIOptions}>
+      {link && <Component url={link.url} text={link.text} range={link.range} />}
+    </GenericPopover>
   );
 };
diff --git a/packages/react/src/components/LinkToolbar/LinkToolbarProps.ts b/packages/react/src/components/LinkToolbar/LinkToolbarProps.ts
index d56fb3d327..cb319e4f6a 100644
--- a/packages/react/src/components/LinkToolbar/LinkToolbarProps.ts
+++ b/packages/react/src/components/LinkToolbar/LinkToolbarProps.ts
@@ -1,18 +1,9 @@
-import {
-  BlockNoteEditor,
-  BlockSchema,
-  LinkToolbarState,
-  InlineContentSchema,
-  StyleSchema,
-  UiElementPosition,
-} from "@blocknote/core";
+import { Range } from "@tiptap/core";
+import { ReactNode } from "react";
 
-export type LinkToolbarProps = Omit<LinkToolbarState, keyof UiElementPosition> &
-  Pick<
-    BlockNoteEditor<
-      BlockSchema,
-      InlineContentSchema,
-      StyleSchema
-    >["linkToolbar"],
-    "deleteLink" | "editLink" | "startHideTimer" | "stopHideTimer"
-  >;
+export type LinkToolbarProps = {
+  url: string;
+  text: string;
+  range: Range;
+  children?: ReactNode;
+};
diff --git a/packages/react/src/components/Popovers/BlockPopover.tsx b/packages/react/src/components/Popovers/BlockPopover.tsx
new file mode 100644
index 0000000000..1262b567f8
--- /dev/null
+++ b/packages/react/src/components/Popovers/BlockPopover.tsx
@@ -0,0 +1,51 @@
+import { getNodeById } from "@blocknote/core";
+import { ReactNode, useMemo } from "react";
+
+import { useBlockNoteEditor } from "../../hooks/useBlockNoteEditor.js";
+import { FloatingUIOptions } from "./FloatingUIOptions.js";
+import { GenericPopover } from "./GenericPopover.js";
+
+export const BlockPopover = (
+  props: FloatingUIOptions & {
+    blockId: string | undefined;
+    children: ReactNode;
+  },
+) => {
+  const { blockId, children, ...rest } = 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(
+    () =>
+      editor.transact((tr) => {
+        if (!blockId) {
+          return undefined;
+        }
+
+        // TODO use the location API for this
+        const nodePosInfo = getNodeById(blockId, tr.doc);
+        if (!nodePosInfo) {
+          return undefined;
+        }
+
+        const { node } = editor.prosemirrorView.domAtPos(
+          nodePosInfo.posBeforeNode + 1,
+        );
+        if (!(node instanceof Element)) {
+          return undefined;
+        }
+
+        return node;
+      }),
+    [editor, blockId],
+  );
+
+  return (
+    <GenericPopover reference={element} {...rest}>
+      {children}
+    </GenericPopover>
+  );
+};
diff --git a/packages/react/src/components/Popovers/FloatingUIOptions.ts b/packages/react/src/components/Popovers/FloatingUIOptions.ts
new file mode 100644
index 0000000000..2d11726428
--- /dev/null
+++ b/packages/react/src/components/Popovers/FloatingUIOptions.ts
@@ -0,0 +1,17 @@
+import {
+  UseDismissProps,
+  UseFloatingOptions,
+  UseHoverProps,
+  UseTransitionStatusProps,
+  UseTransitionStylesProps,
+} from "@floating-ui/react";
+import { HTMLAttributes } from "react";
+
+export type FloatingUIOptions = {
+  useFloatingOptions?: UseFloatingOptions;
+  useTransitionStylesProps?: UseTransitionStylesProps;
+  useTransitionStatusProps?: UseTransitionStatusProps;
+  useDismissProps?: UseDismissProps;
+  useHoverProps?: UseHoverProps;
+  elementProps?: HTMLAttributes<HTMLDivElement>;
+};
diff --git a/packages/react/src/components/Popovers/GenericPopover.tsx b/packages/react/src/components/Popovers/GenericPopover.tsx
new file mode 100644
index 0000000000..0411251d81
--- /dev/null
+++ b/packages/react/src/components/Popovers/GenericPopover.tsx
@@ -0,0 +1,114 @@
+import {
+  useFloating,
+  useTransitionStyles,
+  useDismiss,
+  useInteractions,
+  useMergeRefs,
+  useTransitionStatus,
+  autoUpdate,
+  ReferenceType,
+  useHover,
+} from "@floating-ui/react";
+import { HTMLAttributes, ReactNode, useEffect, useRef } from "react";
+
+import { FloatingUIOptions } from "./FloatingUIOptions.js";
+
+export const GenericPopover = (
+  props: FloatingUIOptions & {
+    reference?: ReferenceType | null;
+    children: ReactNode;
+  },
+) => {
+  const { refs, floatingStyles, context } = useFloating<HTMLDivElement>({
+    whileElementsMounted: autoUpdate,
+    ...props.useFloatingOptions,
+  });
+
+  const { isMounted, styles } = useTransitionStyles(
+    context,
+    props.useTransitionStylesProps,
+  );
+  const { status } = useTransitionStatus(
+    context,
+    props.useTransitionStatusProps,
+  );
+
+  const dismiss = useDismiss(context, props.useDismissProps);
+  const hover = useHover(context, { enabled: false, ...props.useHoverProps });
+  const { getFloatingProps } = useInteractions([dismiss, hover]);
+
+  const innerHTML = useRef<string>("");
+  const ref = useRef<HTMLDivElement>(null);
+  const mergedRefs = useMergeRefs([ref, refs.setFloating]);
+
+  useEffect(() => {
+    if (props.reference) {
+      if (props.reference instanceof Element) {
+        refs.setReference(props.reference);
+      } else {
+        // FloatingUI's virtual elements can have a `contextElement`, which is
+        // useful for automatically calling `getBoundingClientRect` when said
+        // element changes position. However, it's pretty limited as FloatingUI
+        // doesn't attach listeners to `contextElement` for e.g. `useHover`.
+        // Therefore, we also set the reference element to `contextElement` so
+        // event listeners can be attached.
+        if (props.reference.contextElement) {
+          refs.setReference(props.reference.contextElement);
+        }
+        refs.setPositionReference(props.reference);
+      }
+    }
+  }, [props.reference, refs]);
+
+  // Stores the last rendered `innerHTML` of the popover while it was
+  // open. The `innerHTML` is used while the popover is closing, as the
+  // React children may rerender during this time, causing unwanted
+  // behaviour.
+  useEffect(() => {
+    if (status === "initial" || status === "open") {
+      if (ref.current?.innerHTML) {
+        innerHTML.current = ref.current.innerHTML;
+      }
+    }
+  }, [status, props.reference]);
+
+  if (!isMounted) {
+    return false;
+  }
+
+  const mergedProps: HTMLAttributes<HTMLDivElement> = {
+    ...props.elementProps,
+    style: {
+      display: "flex",
+      ...props.elementProps?.style,
+      zIndex: `calc(var(--bn-ui-base-z-index) + ${props.elementProps?.style?.zIndex || 0})`,
+      ...floatingStyles,
+      ...styles,
+    },
+    ...getFloatingProps(),
+  };
+
+  if (status === "close") {
+    // While the popover is closing, shows its last rendered `innerHTML` while
+    // it was open, instead of the React children. This is because they may
+    // rerender during this time, causing unwanted behaviour.
+    //
+    // When we use the `GenericPopover` for BlockNote's internal UI elements
+    // this isn't a huge deal, as we only pass child components if the popover
+    // should be open. So without this fix, the popover just won't transition
+    // out and will instead appear to hide instantly.
+    return (
+      <div
+        ref={mergedRefs}
+        {...mergedProps}
+        dangerouslySetInnerHTML={{ __html: innerHTML.current }}
+      />
+    );
+  }
+
+  return (
+    <div ref={mergedRefs} {...mergedProps}>
+      {props.children}
+    </div>
+  );
+};
diff --git a/packages/react/src/components/Popovers/PositionPopover.tsx b/packages/react/src/components/Popovers/PositionPopover.tsx
new file mode 100644
index 0000000000..22df3d7fb7
--- /dev/null
+++ b/packages/react/src/components/Popovers/PositionPopover.tsx
@@ -0,0 +1,54 @@
+import { posToDOMRect } from "@tiptap/core";
+import { ReactNode, useMemo, useRef } from "react";
+
+import { useBlockNoteEditor } from "../../hooks/useBlockNoteEditor.js";
+import { FloatingUIOptions } from "./FloatingUIOptions.js";
+import { GenericPopover } from "./GenericPopover.js";
+
+export const PositionPopover = (
+  props: FloatingUIOptions & {
+    position: { from: number; to?: number } | undefined;
+    children: ReactNode;
+  },
+) => {
+  const { position, children, ...rest } = 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());
+  const virtualElement = useMemo(
+    () => ({
+      getBoundingClientRect: () => {
+        if (from === undefined || to === undefined) {
+          return boundingClientRect.current;
+        }
+
+        boundingClientRect.current = posToDOMRect(
+          editor.prosemirrorView,
+          from,
+          to ?? from,
+        );
+
+        return boundingClientRect.current;
+      },
+      contextElement:
+        // Use first child as the editor DOM element may itself be scrollable.
+        // For FloatingUI to auto-update the position during scrolling, the
+        // `contextElement` must be a descendant of the scroll container.
+        !headless && editor.domElement?.firstChild instanceof Element
+          ? editor.domElement.firstChild
+          : undefined,
+    }),
+    [editor, headless, from, to],
+  );
+
+  return (
+    <GenericPopover reference={virtualElement} {...rest}>
+      {children}
+    </GenericPopover>
+  );
+};
diff --git a/packages/react/src/components/Popovers/TableCellPopover.tsx b/packages/react/src/components/Popovers/TableCellPopover.tsx
new file mode 100644
index 0000000000..8704169305
--- /dev/null
+++ b/packages/react/src/components/Popovers/TableCellPopover.tsx
@@ -0,0 +1,57 @@
+import { getNodeById } from "@blocknote/core";
+import { ReactNode, useMemo } from "react";
+
+import { useBlockNoteEditor } from "../../hooks/useBlockNoteEditor.js";
+import { FloatingUIOptions } from "./FloatingUIOptions.js";
+import { GenericPopover } from "./GenericPopover.js";
+
+export const TableCellPopover = (
+  props: FloatingUIOptions & {
+    blockId: string | undefined;
+    colIndex: number | undefined;
+    rowIndex: number | undefined;
+    children: ReactNode;
+  },
+) => {
+  const { blockId, colIndex, rowIndex, children, ...rest } = props;
+
+  const editor = useBlockNoteEditor<any, any, any>();
+
+  const element = useMemo(() => {
+    if (
+      blockId === undefined ||
+      colIndex === undefined ||
+      rowIndex === undefined
+    ) {
+      return undefined;
+    }
+
+    // TODO use the location API for this
+    const nodePosInfo = getNodeById(blockId, editor.prosemirrorState.doc);
+    if (!nodePosInfo) {
+      return undefined;
+    }
+
+    const tableBeforePos = nodePosInfo.posBeforeNode + 1;
+
+    const rowBeforePos = editor.prosemirrorState.doc
+      .resolve(tableBeforePos + 1)
+      .posAtIndex(rowIndex || 0);
+    const cellBeforePos = editor.prosemirrorState.doc
+      .resolve(rowBeforePos + 1)
+      .posAtIndex(colIndex || 0);
+
+    const { node } = editor.prosemirrorView.domAtPos(cellBeforePos + 1);
+    if (!(node instanceof Element)) {
+      return undefined;
+    }
+
+    return node;
+  }, [editor, blockId, colIndex, rowIndex]);
+
+  return (
+    <GenericPopover reference={element} {...rest}>
+      {children}
+    </GenericPopover>
+  );
+};
diff --git a/packages/react/src/components/SideMenu/DefaultButtons/AddBlockButton.tsx b/packages/react/src/components/SideMenu/DefaultButtons/AddBlockButton.tsx
index c5a23c2796..56c1d84546 100644
--- a/packages/react/src/components/SideMenu/DefaultButtons/AddBlockButton.tsx
+++ b/packages/react/src/components/SideMenu/DefaultButtons/AddBlockButton.tsx
@@ -1,51 +1,55 @@
-import {
-  BlockSchema,
-  DefaultBlockSchema,
-  DefaultInlineContentSchema,
-  DefaultStyleSchema,
-  InlineContentSchema,
-  StyleSchema,
-} from "@blocknote/core";
+import { SideMenu, SuggestionMenu } from "@blocknote/core/extensions";
 import { AiOutlinePlus } from "react-icons/ai";
 
 import { useCallback } from "react";
 import { useComponentsContext } from "../../../editor/ComponentsContext.js";
 import { useBlockNoteEditor } from "../../../hooks/useBlockNoteEditor.js";
 import { useDictionary } from "../../../i18n/dictionary.js";
-import { SideMenuProps } from "../SideMenuProps.js";
-
-export const AddBlockButton = <
-  BSchema extends BlockSchema = DefaultBlockSchema,
-  I extends InlineContentSchema = DefaultInlineContentSchema,
-  S extends StyleSchema = DefaultStyleSchema,
->(
-  props: Pick<SideMenuProps<BSchema, I, S>, "block">,
-) => {
+import {
+  useExtension,
+  useExtensionState,
+} from "../../../hooks/useExtension.js";
+
+export const AddBlockButton = () => {
   const Components = useComponentsContext()!;
   const dict = useDictionary();
 
-  const editor = useBlockNoteEditor<BSchema, I, S>();
+  const editor = useBlockNoteEditor<any, any, any>();
+
+  const suggestionMenu = useExtension(SuggestionMenu);
+  const block = useExtensionState(SideMenu, {
+    editor,
+    selector: (state) => state?.block,
+  });
 
   const onClick = useCallback(() => {
-    const blockContent = props.block.content;
+    if (block === undefined) {
+      return;
+    }
+
+    const blockContent = block.content;
     const isBlockEmpty =
       blockContent !== undefined &&
       Array.isArray(blockContent) &&
       blockContent.length === 0;
 
     if (isBlockEmpty) {
-      editor.setTextCursorPosition(props.block);
-      editor.openSuggestionMenu("/");
+      editor.setTextCursorPosition(block);
+      suggestionMenu.openSuggestionMenu("/");
     } else {
       const insertedBlock = editor.insertBlocks(
         [{ type: "paragraph" }],
-        props.block,
+        block,
         "after",
       )[0];
       editor.setTextCursorPosition(insertedBlock);
-      editor.openSuggestionMenu("/");
+      suggestionMenu.openSuggestionMenu("/");
     }
-  }, [editor, props.block]);
+  }, [block, editor, suggestionMenu]);
+
+  if (block === undefined) {
+    return null;
+  }
 
   return (
     <Components.SideMenu.Button
diff --git a/packages/react/src/components/SideMenu/DefaultButtons/DragHandleButton.tsx b/packages/react/src/components/SideMenu/DefaultButtons/DragHandleButton.tsx
index 82f1ac7310..86bd065192 100644
--- a/packages/react/src/components/SideMenu/DefaultButtons/DragHandleButton.tsx
+++ b/packages/react/src/components/SideMenu/DefaultButtons/DragHandleButton.tsx
@@ -1,42 +1,41 @@
-import {
-  BlockSchema,
-  DefaultBlockSchema,
-  DefaultInlineContentSchema,
-  DefaultStyleSchema,
-  InlineContentSchema,
-  StyleSchema,
-} from "@blocknote/core";
+import { SideMenu } from "@blocknote/core/extensions";
 import { MdDragIndicator } from "react-icons/md";
 
 import { useComponentsContext } from "../../../editor/ComponentsContext.js";
 import { useDictionary } from "../../../i18n/dictionary.js";
 import { DragHandleMenu } from "../DragHandleMenu/DragHandleMenu.js";
 import { SideMenuProps } from "../SideMenuProps.js";
+import {
+  useExtension,
+  useExtensionState,
+} from "../../../hooks/useExtension.js";
 
-export const DragHandleButton = <
-  BSchema extends BlockSchema = DefaultBlockSchema,
-  I extends InlineContentSchema = DefaultInlineContentSchema,
-  S extends StyleSchema = DefaultStyleSchema,
->(
-  props: Omit<SideMenuProps<BSchema, I, S>, "addBlock"> & {
-    /**
-     * The menu items to render.
-     */
+export const DragHandleButton = (
+  props: SideMenuProps & {
     children?: React.ReactNode;
   },
 ) => {
   const Components = useComponentsContext()!;
   const dict = useDictionary();
 
+  const sideMenu = useExtension(SideMenu);
+  const block = useExtensionState(SideMenu, {
+    selector: (state) => state?.block,
+  });
+
+  if (block === undefined) {
+    return null;
+  }
+
   const Component = props.dragHandleMenu || DragHandleMenu;
 
   return (
     <Components.Generic.Menu.Root
       onOpenChange={(open: boolean) => {
         if (open) {
-          props.freezeMenu();
+          sideMenu.freezeMenu();
         } else {
-          props.unfreezeMenu();
+          sideMenu.unfreezeMenu();
         }
       }}
       position={"left"}
@@ -45,13 +44,13 @@ export const DragHandleButton = <
         <Components.SideMenu.Button
           label={dict.side_menu.drag_handle_label}
           draggable={true}
-          onDragStart={(e) => props.blockDragStart(e, props.block)}
-          onDragEnd={props.blockDragEnd}
+          onDragStart={(e) => sideMenu.blockDragStart(e, block)}
+          onDragEnd={sideMenu.blockDragEnd}
           className={"bn-button"}
           icon={<MdDragIndicator size={24} data-test="dragHandle" />}
         />
       </Components.Generic.Menu.Trigger>
-      <Component block={props.block}>{props.children}</Component>
+      <Component>{props.children}</Component>
     </Components.Generic.Menu.Root>
   );
 };
diff --git a/packages/react/src/components/SideMenu/DragHandleMenu/DefaultItems/BlockColorsItem.tsx b/packages/react/src/components/SideMenu/DragHandleMenu/DefaultItems/BlockColorsItem.tsx
index 24a738960e..df84a51fbb 100644
--- a/packages/react/src/components/SideMenu/DragHandleMenu/DefaultItems/BlockColorsItem.tsx
+++ b/packages/react/src/components/SideMenu/DragHandleMenu/DefaultItems/BlockColorsItem.tsx
@@ -1,46 +1,30 @@
-import {
-  Block,
-  blockHasType,
-  BlockSchema,
-  DefaultBlockSchema,
-  DefaultInlineContentSchema,
-  DefaultStyleSchema,
-  editorHasBlockWithType,
-  InlineContentSchema,
-  StyleSchema,
-} from "@blocknote/core";
+import { blockHasType, editorHasBlockWithType } from "@blocknote/core";
+import { SideMenu } from "@blocknote/core/extensions";
 import { ReactNode } from "react";
 
 import { useComponentsContext } from "../../../../editor/ComponentsContext.js";
 import { useBlockNoteEditor } from "../../../../hooks/useBlockNoteEditor.js";
 import { ColorPicker } from "../../../ColorPicker/ColorPicker.js";
-import { DragHandleMenuProps } from "../DragHandleMenuProps.js";
+import { useExtensionState } from "../../../../hooks/useExtension.js";
 
-export const BlockColorsItem = <
-  BSchema extends BlockSchema = DefaultBlockSchema,
-  I extends InlineContentSchema = DefaultInlineContentSchema,
-  S extends StyleSchema = DefaultStyleSchema,
->(
-  props: DragHandleMenuProps<BSchema, I, S> & {
-    children: ReactNode;
-  },
-) => {
+export const BlockColorsItem = (props: { children: ReactNode }) => {
   const Components = useComponentsContext()!;
 
-  const editor = useBlockNoteEditor<BSchema, I, S>();
+  const editor = useBlockNoteEditor<any, any, any>();
 
-  // We cast the block to a generic one, as the base type causes type errors
-  // with runtime type checking using `blockHasType`. Runtime type checking is
-  // more valuable than static checks, so better to do it like this.
-  const block = props.block as Block<any, any, any>;
+  const block = useExtensionState(SideMenu, {
+    editor,
+    selector: (state) => state?.block,
+  });
 
   if (
-    !blockHasType(block, editor, block.type, {
+    block === undefined ||
+    (!blockHasType(block, editor, block.type, {
       textColor: "string",
     }) &&
-    !blockHasType(block, editor, block.type, {
-      backgroundColor: "string",
-    })
+      !blockHasType(block, editor, block.type, {
+        backgroundColor: "string",
+      }))
   ) {
     return null;
   }
diff --git a/packages/react/src/components/SideMenu/DragHandleMenu/DefaultItems/RemoveBlockItem.tsx b/packages/react/src/components/SideMenu/DragHandleMenu/DefaultItems/RemoveBlockItem.tsx
index 83bb1ece34..89328fb029 100644
--- a/packages/react/src/components/SideMenu/DragHandleMenu/DefaultItems/RemoveBlockItem.tsx
+++ b/packages/react/src/components/SideMenu/DragHandleMenu/DefaultItems/RemoveBlockItem.tsx
@@ -1,34 +1,28 @@
-import {
-  BlockSchema,
-  DefaultBlockSchema,
-  DefaultInlineContentSchema,
-  DefaultStyleSchema,
-  InlineContentSchema,
-  StyleSchema,
-} from "@blocknote/core";
+import { SideMenu } from "@blocknote/core/extensions";
 import { ReactNode } from "react";
 
 import { useComponentsContext } from "../../../../editor/ComponentsContext.js";
 import { useBlockNoteEditor } from "../../../../hooks/useBlockNoteEditor.js";
-import { DragHandleMenuProps } from "../DragHandleMenuProps.js";
+import { useExtensionState } from "../../../../hooks/useExtension.js";
 
-export const RemoveBlockItem = <
-  BSchema extends BlockSchema = DefaultBlockSchema,
-  I extends InlineContentSchema = DefaultInlineContentSchema,
-  S extends StyleSchema = DefaultStyleSchema,
->(
-  props: DragHandleMenuProps<BSchema, I, S> & {
-    children: ReactNode;
-  },
-) => {
+export const RemoveBlockItem = (props: { children: ReactNode }) => {
   const Components = useComponentsContext()!;
 
-  const editor = useBlockNoteEditor<BSchema, I, S>();
+  const editor = useBlockNoteEditor<any, any, any>();
+
+  const block = useExtensionState(SideMenu, {
+    editor,
+    selector: (state) => state?.block,
+  });
+
+  if (block === undefined) {
+    return null;
+  }
 
   return (
     <Components.Generic.Menu.Item
       className={"bn-menu-item"}
-      onClick={() => editor.removeBlocks([props.block])}
+      onClick={() => editor.removeBlocks([block])}
     >
       {props.children}
     </Components.Generic.Menu.Item>
diff --git a/packages/react/src/components/SideMenu/DragHandleMenu/DefaultItems/TableHeadersItem.tsx b/packages/react/src/components/SideMenu/DragHandleMenu/DefaultItems/TableHeadersItem.tsx
index 29917f00ac..feefce2d79 100644
--- a/packages/react/src/components/SideMenu/DragHandleMenu/DefaultItems/TableHeadersItem.tsx
+++ b/packages/react/src/components/SideMenu/DragHandleMenu/DefaultItems/TableHeadersItem.tsx
@@ -1,55 +1,50 @@
-import {
-  BlockSchema,
-  DefaultBlockSchema,
-  DefaultInlineContentSchema,
-  DefaultStyleSchema,
-  InlineContentSchema,
-  SpecificBlock,
-  StyleSchema,
-} from "@blocknote/core";
+import { DefaultBlockSchema, SpecificBlock } from "@blocknote/core";
+import { SideMenu } from "@blocknote/core/extensions";
 import { ReactNode } from "react";
 
 import { useComponentsContext } from "../../../../editor/ComponentsContext.js";
 import { useBlockNoteEditor } from "../../../../hooks/useBlockNoteEditor.js";
-import { DragHandleMenuProps } from "../DragHandleMenuProps.js";
+import { useExtensionState } from "../../../../hooks/useExtension.js";
 
-export const TableRowHeaderItem = <
-  BSchema extends BlockSchema = DefaultBlockSchema,
-  I extends InlineContentSchema = DefaultInlineContentSchema,
-  S extends StyleSchema = DefaultStyleSchema,
->(
-  props: Omit<DragHandleMenuProps<BSchema, I, S>, "block"> & {
-    block: SpecificBlock<{ table: DefaultBlockSchema["table"] }, "table", I, S>;
-    children: ReactNode;
-  },
-) => {
+export const TableRowHeaderItem = (props: { children: ReactNode }) => {
   const Components = useComponentsContext()!;
 
   const editor = useBlockNoteEditor<
     { table: DefaultBlockSchema["table"] },
-    I,
-    S
+    any,
+    any
   >();
 
-  if (props.block.type !== "table" || !editor.settings.tables.headers) {
+  const block = useExtensionState(SideMenu, {
+    editor,
+    selector: (state) =>
+      state?.block as
+        | SpecificBlock<
+            { table: DefaultBlockSchema["table"] },
+            "table",
+            any,
+            any
+          >
+        | undefined,
+  });
+
+  if (
+    block === undefined ||
+    block.type !== "table" ||
+    !editor.settings.tables.headers
+  ) {
     return null;
   }
 
   // TODO only support 1 header row for now
-  const isHeaderRow = Boolean(props.block.content.headerRows);
+  const isHeaderRow = Boolean(block.content.headerRows);
 
   return (
     <Components.Generic.Menu.Item
       className={"bn-menu-item"}
       checked={isHeaderRow}
       onClick={() => {
-        // The block may have been modified and out of date, so we get the latest block
-        const block = editor.getBlock(props.block.id);
-        if (!block) {
-          return;
-        }
         editor.updateBlock(block, {
-          ...block,
           content: {
             ...block.content,
             headerRows: isHeaderRow ? undefined : 1,
@@ -62,41 +57,47 @@ export const TableRowHeaderItem = <
   );
 };
 
-export const TableColumnHeaderItem = <
-  BSchema extends BlockSchema = DefaultBlockSchema,
-  I extends InlineContentSchema = DefaultInlineContentSchema,
-  S extends StyleSchema = DefaultStyleSchema,
->(
-  props: Omit<DragHandleMenuProps<BSchema, I, S>, "block"> & {
-    block: SpecificBlock<{ table: DefaultBlockSchema["table"] }, "table", I, S>;
-    children: ReactNode;
-  },
-) => {
+export const TableColumnHeaderItem = (props: { children: ReactNode }) => {
   const Components = useComponentsContext()!;
 
   const editor = useBlockNoteEditor<
     { table: DefaultBlockSchema["table"] },
-    I,
-    S
+    any,
+    any
   >();
 
-  if (props.block.type !== "table" || !editor.settings.tables.headers) {
+  const block = useExtensionState(SideMenu, {
+    editor,
+    selector: (state) =>
+      state?.block as
+        | SpecificBlock<
+            { table: DefaultBlockSchema["table"] },
+            "table",
+            any,
+            any
+          >
+        | undefined,
+  });
+
+  if (
+    block === undefined ||
+    block.type !== "table" ||
+    !editor.settings.tables.headers
+  ) {
     return null;
   }
 
   // TODO only support 1 header column for now
-  const isHeaderColumn = Boolean(props.block.content.headerCols);
+  const isHeaderColumn = Boolean(block.content.headerCols);
 
   return (
     <Components.Generic.Menu.Item
       className={"bn-menu-item"}
       checked={isHeaderColumn}
       onClick={() => {
-        editor.updateBlock(props.block, {
-          type: "table",
+        editor.updateBlock(block, {
           content: {
-            ...props.block.content,
-            type: "tableContent",
+            ...block.content,
             headerCols: isHeaderColumn ? undefined : 1,
           },
         });
diff --git a/packages/react/src/components/SideMenu/DragHandleMenu/DragHandleMenu.tsx b/packages/react/src/components/SideMenu/DragHandleMenu/DragHandleMenu.tsx
index 6565992e92..2a443eb40c 100644
--- a/packages/react/src/components/SideMenu/DragHandleMenu/DragHandleMenu.tsx
+++ b/packages/react/src/components/SideMenu/DragHandleMenu/DragHandleMenu.tsx
@@ -1,18 +1,9 @@
-import {
-  BlockSchema,
-  DefaultBlockSchema,
-  DefaultInlineContentSchema,
-  DefaultStyleSchema,
-  InlineContentSchema,
-  StyleSchema,
-} from "@blocknote/core";
 import { ReactNode } from "react";
 
 import { useComponentsContext } from "../../../editor/ComponentsContext.js";
 import { useDictionary } from "../../../i18n/dictionary.js";
 import { BlockColorsItem } from "./DefaultItems/BlockColorsItem.js";
 import { RemoveBlockItem } from "./DefaultItems/RemoveBlockItem.js";
-import { DragHandleMenuProps } from "./DragHandleMenuProps.js";
 import {
   TableColumnHeaderItem,
   TableRowHeaderItem,
@@ -26,13 +17,7 @@ import {
  * - Default items: Components found within the `/DefaultItems` directory.
  * - Custom items: The `DragHandleMenuItem` component.
  */
-export const DragHandleMenu = <
-  BSchema extends BlockSchema = DefaultBlockSchema,
-  I extends InlineContentSchema = DefaultInlineContentSchema,
-  S extends StyleSchema = DefaultStyleSchema,
->(
-  props: DragHandleMenuProps<BSchema, I, S> & { children?: ReactNode },
-) => {
+export const DragHandleMenu = (props: { children?: ReactNode }) => {
   const Components = useComponentsContext()!;
   const dict = useDictionary();
 
@@ -42,16 +27,12 @@ export const DragHandleMenu = <
     >
       {props.children || (
         <>
-          <RemoveBlockItem {...props}>
-            {dict.drag_handle.delete_menuitem}
-          </RemoveBlockItem>
-          <BlockColorsItem {...props}>
-            {dict.drag_handle.colors_menuitem}
-          </BlockColorsItem>
-          <TableRowHeaderItem {...(props as any)}>
+          <RemoveBlockItem>{dict.drag_handle.delete_menuitem}</RemoveBlockItem>
+          <BlockColorsItem>{dict.drag_handle.colors_menuitem}</BlockColorsItem>
+          <TableRowHeaderItem>
             {dict.drag_handle.header_row_menuitem}
           </TableRowHeaderItem>
-          <TableColumnHeaderItem {...(props as any)}>
+          <TableColumnHeaderItem>
             {dict.drag_handle.header_column_menuitem}
           </TableColumnHeaderItem>
         </>
diff --git a/packages/react/src/components/SideMenu/DragHandleMenu/DragHandleMenuProps.ts b/packages/react/src/components/SideMenu/DragHandleMenu/DragHandleMenuProps.ts
deleted file mode 100644
index 714d2c09d3..0000000000
--- a/packages/react/src/components/SideMenu/DragHandleMenu/DragHandleMenuProps.ts
+++ /dev/null
@@ -1,17 +0,0 @@
-import {
-  Block,
-  BlockSchema,
-  DefaultBlockSchema,
-  DefaultInlineContentSchema,
-  DefaultStyleSchema,
-  InlineContentSchema,
-  StyleSchema,
-} from "@blocknote/core";
-
-export type DragHandleMenuProps<
-  BSchema extends BlockSchema = DefaultBlockSchema,
-  I extends InlineContentSchema = DefaultInlineContentSchema,
-  S extends StyleSchema = DefaultStyleSchema,
-> = {
-  block: Block<BSchema, I, S>;
-};
diff --git a/packages/react/src/components/SideMenu/SideMenu.tsx b/packages/react/src/components/SideMenu/SideMenu.tsx
index 64eda02c05..7a11f5e1c9 100644
--- a/packages/react/src/components/SideMenu/SideMenu.tsx
+++ b/packages/react/src/components/SideMenu/SideMenu.tsx
@@ -1,14 +1,9 @@
-import {
-  BlockSchema,
-  DefaultBlockSchema,
-  DefaultInlineContentSchema,
-  DefaultStyleSchema,
-  InlineContentSchema,
-  StyleSchema,
-} from "@blocknote/core";
+import { SideMenu as SideMenuExtension } from "@blocknote/core/extensions";
 import { ReactNode, useMemo } from "react";
 
 import { useComponentsContext } from "../../editor/ComponentsContext.js";
+import { useBlockNoteEditor } from "../../hooks/useBlockNoteEditor.js";
+import { useExtensionState } from "../../hooks/useExtension.js";
 import { AddBlockButton } from "./DefaultButtons/AddBlockButton.js";
 import { DragHandleButton } from "./DefaultButtons/DragHandleButton.js";
 import { SideMenuProps } from "./SideMenuProps.js";
@@ -22,29 +17,33 @@ import { SideMenuProps } from "./SideMenuProps.js";
  * - Default buttons: Components found within the `/DefaultButtons` directory.
  * - Custom buttons: The `SideMenuButton` component.
  */
-export const SideMenu = <
-  BSchema extends BlockSchema = DefaultBlockSchema,
-  I extends InlineContentSchema = DefaultInlineContentSchema,
-  S extends StyleSchema = DefaultStyleSchema,
->(
-  props: SideMenuProps<BSchema, I, S> & { children?: ReactNode },
-) => {
+export const SideMenu = (props: SideMenuProps & { children?: ReactNode }) => {
   const Components = useComponentsContext()!;
 
+  const editor = useBlockNoteEditor<any, any, any>();
+
+  const block = useExtensionState(SideMenuExtension, {
+    editor,
+    selector: (state) => state?.block,
+  });
+
   const dataAttributes = useMemo(() => {
+    if (block === undefined) {
+      return {};
+    }
+
     const attrs: Record<string, string> = {
-      "data-block-type": props.block.type,
+      "data-block-type": block.type,
     };
 
-    if (props.block.type === "heading") {
-      attrs["data-level"] = (props.block.props as any).level.toString();
+    if (block.type === "heading") {
+      attrs["data-level"] = (block.props as any).level.toString();
     }
 
     if (
-      props.editor.schema.blockSpecs[props.block.type].implementation.meta
-        ?.fileBlockAccept
+      editor.schema.blockSpecs[block.type].implementation.meta?.fileBlockAccept
     ) {
-      if (props.block.props.url) {
+      if (block.props.url) {
         attrs["data-url"] = "true";
       } else {
         attrs["data-url"] = "false";
@@ -52,14 +51,14 @@ export const SideMenu = <
     }
 
     return attrs;
-  }, [props.block.props, props.block.type, props.editor.schema.blockSpecs]);
+  }, [block, editor.schema.blockSpecs]);
 
   return (
     <Components.SideMenu.Root className={"bn-side-menu"} {...dataAttributes}>
       {props.children || (
         <>
-          <AddBlockButton {...props} />
-          <DragHandleButton {...props} />
+          <AddBlockButton />
+          <DragHandleButton />
         </>
       )}
     </Components.SideMenu.Root>
diff --git a/packages/react/src/components/SideMenu/SideMenuController.tsx b/packages/react/src/components/SideMenu/SideMenuController.tsx
index 935dd9865b..be27a51c2e 100644
--- a/packages/react/src/components/SideMenu/SideMenuController.tsx
+++ b/packages/react/src/components/SideMenu/SideMenuController.tsx
@@ -1,61 +1,53 @@
-import {
-  BlockSchema,
-  DefaultBlockSchema,
-  DefaultInlineContentSchema,
-  DefaultStyleSchema,
-  InlineContentSchema,
-  StyleSchema,
-} from "@blocknote/core";
-import { FC } from "react";
+import { SideMenu as SideMenuExtension } from "@blocknote/core/extensions";
+import { FC, useMemo } from "react";
 
-import { UseFloatingOptions } from "@floating-ui/react";
-import { useBlockNoteEditor } from "../../hooks/useBlockNoteEditor.js";
-import { useUIElementPositioning } from "../../hooks/useUIElementPositioning.js";
-import { useUIPluginState } from "../../hooks/useUIPluginState.js";
+import { useExtensionState } from "../../hooks/useExtension.js";
+import { BlockPopover } from "../Popovers/BlockPopover.js";
+import { FloatingUIOptions } from "../Popovers/FloatingUIOptions.js";
 import { SideMenu } from "./SideMenu.js";
 import { SideMenuProps } from "./SideMenuProps.js";
 
-export const SideMenuController = <
-  BSchema extends BlockSchema = DefaultBlockSchema,
-  I extends InlineContentSchema = DefaultInlineContentSchema,
-  S extends StyleSchema = DefaultStyleSchema,
->(props: {
-  sideMenu?: FC<SideMenuProps<BSchema, I, S>>;
-  floatingOptions?: Partial<UseFloatingOptions>;
+export const SideMenuController = (props: {
+  sideMenu?: FC<SideMenuProps>;
+  floatingUIOptions?: Partial<FloatingUIOptions>;
 }) => {
-  const editor = useBlockNoteEditor<BSchema, I, S>();
-
-  const callbacks = {
-    blockDragStart: editor.sideMenu.blockDragStart,
-    blockDragEnd: editor.sideMenu.blockDragEnd,
-    freezeMenu: editor.sideMenu.freezeMenu,
-    unfreezeMenu: editor.sideMenu.unfreezeMenu,
-  };
-
-  const state = useUIPluginState(
-    editor.sideMenu.onUpdate.bind(editor.sideMenu),
-  );
-  const { isMounted, ref, style, getFloatingProps } = useUIElementPositioning(
-    state?.show || false,
-    state?.referencePos || null,
-    1000,
-    {
-      placement: "left-start",
-      ...props.floatingOptions,
+  const state = useExtensionState(SideMenuExtension, {
+    selector: (state) => {
+      return state !== undefined
+        ? {
+            show: state.show,
+            block: state.block,
+          }
+        : undefined;
     },
+  });
+
+  const { show, block } = state || {};
+
+  const floatingUIOptions = useMemo<FloatingUIOptions>(
+    () => ({
+      useFloatingOptions: {
+        open: show,
+        placement: "left-start",
+      },
+      useDismissProps: {
+        enabled: false,
+      },
+      elementProps: {
+        style: {
+          zIndex: 20,
+        },
+      },
+      ...props.floatingUIOptions,
+    }),
+    [props.floatingUIOptions, show],
   );
 
-  if (!isMounted || !state) {
-    return null;
-  }
-
-  const { show, referencePos, ...data } = state;
-
   const Component = props.sideMenu || SideMenu;
 
   return (
-    <div ref={ref} style={style} {...getFloatingProps()}>
-      <Component {...data} {...callbacks} editor={editor} />
-    </div>
+    <BlockPopover blockId={block?.id} {...floatingUIOptions}>
+      <Component />
+    </BlockPopover>
   );
 };
diff --git a/packages/react/src/components/SideMenu/SideMenuProps.ts b/packages/react/src/components/SideMenu/SideMenuProps.ts
index a723e3d8ac..329996ee30 100644
--- a/packages/react/src/components/SideMenu/SideMenuProps.ts
+++ b/packages/react/src/components/SideMenu/SideMenuProps.ts
@@ -1,27 +1,5 @@
-import {
-  BlockNoteEditor,
-  BlockSchema,
-  DefaultBlockSchema,
-  DefaultInlineContentSchema,
-  DefaultStyleSchema,
-  InlineContentSchema,
-  SideMenuState,
-  StyleSchema,
-  UiElementPosition,
-} from "@blocknote/core";
 import { FC } from "react";
 
-import { DragHandleMenuProps } from "./DragHandleMenu/DragHandleMenuProps.js";
-
-export type SideMenuProps<
-  BSchema extends BlockSchema = DefaultBlockSchema,
-  I extends InlineContentSchema = DefaultInlineContentSchema,
-  S extends StyleSchema = DefaultStyleSchema,
-> = {
-  editor: BlockNoteEditor<BSchema, I, S>;
-  dragHandleMenu?: FC<DragHandleMenuProps<BSchema, I, S>>;
-} & Omit<SideMenuState<BSchema, I, S>, keyof UiElementPosition> &
-  Pick<
-    BlockNoteEditor<BSchema, I, S>["sideMenu"],
-    "blockDragStart" | "blockDragEnd" | "freezeMenu" | "unfreezeMenu"
-  >;
+export type SideMenuProps = {
+  dragHandleMenu?: FC;
+};
diff --git a/packages/react/src/components/SuggestionMenu/GridSuggestionMenu/GridSuggestionMenuController.tsx b/packages/react/src/components/SuggestionMenu/GridSuggestionMenu/GridSuggestionMenuController.tsx
index ca08083ad7..d712d64f9d 100644
--- a/packages/react/src/components/SuggestionMenu/GridSuggestionMenu/GridSuggestionMenuController.tsx
+++ b/packages/react/src/components/SuggestionMenu/GridSuggestionMenu/GridSuggestionMenuController.tsx
@@ -1,15 +1,15 @@
-import {
-  BlockSchema,
-  InlineContentSchema,
-  StyleSchema,
-  SuggestionMenuState,
-} from "@blocknote/core";
-import { flip, offset, size, UseFloatingOptions } from "@floating-ui/react";
-import { FC, useCallback, useMemo } from "react";
+import { BlockSchema, InlineContentSchema, StyleSchema } from "@blocknote/core";
+import { SuggestionMenu } from "@blocknote/core/extensions";
+import { flip, offset, shift, size, VirtualElement } from "@floating-ui/react";
+import { FC, useEffect, useMemo } from "react";
 
 import { useBlockNoteEditor } from "../../../hooks/useBlockNoteEditor.js";
-import { useUIElementPositioning } from "../../../hooks/useUIElementPositioning.js";
-import { useUIPluginState } from "../../../hooks/useUIPluginState.js";
+import {
+  useExtension,
+  useExtensionState,
+} from "../../../hooks/useExtension.js";
+import { FloatingUIOptions } from "../../Popovers/FloatingUIOptions.js";
+import { GenericPopover } from "../../Popovers/GenericPopover.js";
 import { getDefaultReactEmojiPickerItems } from "./getDefaultReactEmojiPickerItems.js";
 import { GridSuggestionMenu } from "./GridSuggestionMenu.js";
 import { GridSuggestionMenuWrapper } from "./GridSuggestionMenuWrapper.js";
@@ -35,7 +35,7 @@ export function GridSuggestionMenuController<
     getItems?: GetItemsType;
     columns: number;
     minQueryLength?: number;
-    floatingOptions?: Partial<UseFloatingOptions>;
+    floatingUIOptions?: FloatingUIOptions;
   } & (ItemType<GetItemsType> extends DefaultReactGridSuggestionItem
     ? {
         // can be undefined
@@ -65,7 +65,6 @@ export function GridSuggestionMenuController<
     minQueryLength,
     onItemClick,
     getItems,
-    floatingOptions,
   } = props;
 
   const onItemClickOrDefault = useMemo(() => {
@@ -88,52 +87,72 @@ export function GridSuggestionMenuController<
     );
   }, [editor, getItems])!;
 
-  const callbacks = {
-    closeMenu: editor.suggestionMenus.closeMenu,
-    clearQuery: editor.suggestionMenus.clearQuery,
-  };
+  const suggestionMenu = useExtension(SuggestionMenu);
 
-  const cb = useCallback(
-    (callback: (state: SuggestionMenuState) => void) => {
-      return editor.suggestionMenus.onUpdate(triggerCharacter, callback);
-    },
-    [editor.suggestionMenus, triggerCharacter],
-  );
+  useEffect(() => {
+    suggestionMenu.addTriggerCharacter(triggerCharacter);
+  }, [
+    props.triggerCharacter,
+    suggestionMenu,
+    suggestionMenu.store,
+    triggerCharacter,
+  ]);
 
-  const state = useUIPluginState(cb);
+  const state = useExtensionState(SuggestionMenu);
+  const referencePos = useExtensionState(SuggestionMenu, {
+    selector: (state) => state?.referencePos || new DOMRect(),
+  });
 
-  const { isMounted, ref, style, getFloatingProps } = useUIElementPositioning(
-    state?.show || false,
-    state?.referencePos || null,
-    2000,
-    {
-      placement: "bottom-start",
-      middleware: [
-        offset(10),
-        // Flips the menu placement to maximize the space available, and prevents
-        // the menu from being cut off by the confines of the screen.
-        flip(),
-        size({
-          apply({ availableHeight, elements }) {
-            Object.assign(elements.floating.style, {
-              maxHeight: `${availableHeight - 10}px`,
-            });
-          },
-        }),
-      ],
-      onOpenChange(open) {
-        if (!open) {
-          editor.suggestionMenus.closeMenu();
-        }
+  const virtualElement = useMemo<VirtualElement>(
+    () => ({
+      getBoundingClientRect: () => referencePos,
+    }),
+    [referencePos],
+  );
+
+  const floatingUIOptions = useMemo<FloatingUIOptions>(
+    () => ({
+      useFloatingOptions: {
+        open: state?.show && state?.triggerCharacter === triggerCharacter,
+        placement: "bottom-start",
+        middleware: [
+          offset(10),
+          // Flips the menu placement to maximize the space available, and prevents
+          // the menu from being cut off by the confines of the screen.
+          flip({
+            mainAxis: true,
+            crossAxis: false,
+          }),
+          shift(),
+          size({
+            apply({ availableHeight, elements }) {
+              Object.assign(elements.floating.style, {
+                maxHeight: `${availableHeight - 10}px`,
+              });
+            },
+          }),
+        ],
+      },
+      elementProps: {
+        // Prevents editor blurring when clicking the scroll bar.
+        onMouseDownCapture: (event) => event.preventDefault(),
+        style: {
+          zIndex: 70,
+        },
       },
-      ...floatingOptions,
-    },
+      ...props.floatingUIOptions,
+    }),
+    [
+      props.floatingUIOptions,
+      state?.show,
+      state?.triggerCharacter,
+      triggerCharacter,
+    ],
   );
 
   if (
-    !isMounted ||
     !state ||
-    (!state?.ignoreQueryLength &&
+    (!state.ignoreQueryLength &&
       minQueryLength &&
       (state.query.startsWith(" ") || state.query.length < minQueryLength))
   ) {
@@ -141,19 +160,21 @@ export function GridSuggestionMenuController<
   }
 
   return (
-    <div ref={ref} style={style} {...getFloatingProps()}>
-      <GridSuggestionMenuWrapper
-        query={state.query}
-        closeMenu={callbacks.closeMenu}
-        clearQuery={callbacks.clearQuery}
-        getItems={getItemsOrDefault}
-        columns={columns}
-        gridSuggestionMenuComponent={
-          gridSuggestionMenuComponent ||
-          GridSuggestionMenu<ItemType<GetItemsType>>
-        }
-        onItemClick={onItemClickOrDefault}
-      />
-    </div>
+    <GenericPopover reference={virtualElement} {...floatingUIOptions}>
+      {triggerCharacter && (
+        <GridSuggestionMenuWrapper
+          query={state.query}
+          closeMenu={suggestionMenu.closeMenu}
+          clearQuery={suggestionMenu.clearQuery}
+          getItems={getItemsOrDefault}
+          columns={columns}
+          gridSuggestionMenuComponent={
+            gridSuggestionMenuComponent ||
+            GridSuggestionMenu<ItemType<GetItemsType>>
+          }
+          onItemClick={onItemClickOrDefault}
+        />
+      )}
+    </GenericPopover>
   );
 }
diff --git a/packages/react/src/components/SuggestionMenu/GridSuggestionMenu/getDefaultReactEmojiPickerItems.tsx b/packages/react/src/components/SuggestionMenu/GridSuggestionMenu/getDefaultReactEmojiPickerItems.tsx
index b77e41d0cb..de6a629ab0 100644
--- a/packages/react/src/components/SuggestionMenu/GridSuggestionMenu/getDefaultReactEmojiPickerItems.tsx
+++ b/packages/react/src/components/SuggestionMenu/GridSuggestionMenu/getDefaultReactEmojiPickerItems.tsx
@@ -1,10 +1,10 @@
 import {
   BlockNoteEditor,
   BlockSchema,
-  getDefaultEmojiPickerItems,
   InlineContentSchema,
   StyleSchema,
 } from "@blocknote/core";
+import { getDefaultEmojiPickerItems } from "@blocknote/core/extensions";
 import { DefaultReactGridSuggestionItem } from "./types.js";
 
 export async function getDefaultReactEmojiPickerItems<
diff --git a/packages/react/src/components/SuggestionMenu/GridSuggestionMenu/types.tsx b/packages/react/src/components/SuggestionMenu/GridSuggestionMenu/types.tsx
index c24b4f26d4..cb8f9b8f17 100644
--- a/packages/react/src/components/SuggestionMenu/GridSuggestionMenu/types.tsx
+++ b/packages/react/src/components/SuggestionMenu/GridSuggestionMenu/types.tsx
@@ -1,4 +1,4 @@
-import { DefaultGridSuggestionItem } from "@blocknote/core";
+import { DefaultGridSuggestionItem } from "@blocknote/core/extensions";
 import { JSX } from "react";
 
 import { SuggestionMenuProps } from "../types.js";
diff --git a/packages/react/src/components/SuggestionMenu/SuggestionMenuController.tsx b/packages/react/src/components/SuggestionMenu/SuggestionMenuController.tsx
index f904b8169e..cd172928ce 100644
--- a/packages/react/src/components/SuggestionMenu/SuggestionMenuController.tsx
+++ b/packages/react/src/components/SuggestionMenu/SuggestionMenuController.tsx
@@ -1,22 +1,22 @@
+import { BlockSchema, InlineContentSchema, StyleSchema } from "@blocknote/core";
 import {
-  BlockSchema,
-  InlineContentSchema,
-  StyleSchema,
-  SuggestionMenuState,
+  SuggestionMenu as SuggestionMenuExtension,
   filterSuggestionItems,
-} from "@blocknote/core";
+} from "@blocknote/core/extensions";
 import {
   UseFloatingOptions,
+  VirtualElement,
   flip,
   offset,
   shift,
   size,
 } from "@floating-ui/react";
-import { FC, useCallback, useMemo } from "react";
+import { FC, useEffect, useMemo } from "react";
 
 import { useBlockNoteEditor } from "../../hooks/useBlockNoteEditor.js";
-import { useUIElementPositioning } from "../../hooks/useUIElementPositioning.js";
-import { useUIPluginState } from "../../hooks/useUIPluginState.js";
+import { useExtension, useExtensionState } from "../../hooks/useExtension.js";
+import { GenericPopover } from "../Popovers/GenericPopover.js";
+import { FloatingUIOptions } from "../Popovers/FloatingUIOptions.js";
 import { SuggestionMenu } from "./SuggestionMenu.js";
 import { SuggestionMenuWrapper } from "./SuggestionMenuWrapper.js";
 import { getDefaultReactSlashMenuItems } from "./getDefaultReactSlashMenuItems.js";
@@ -38,7 +38,7 @@ export function SuggestionMenuController<
     triggerCharacter: string;
     getItems?: GetItemsType;
     minQueryLength?: number;
-    floatingOptions?: Partial<UseFloatingOptions>;
+    floatingUIOptions?: UseFloatingOptions;
   } & (ItemType<GetItemsType> extends DefaultReactSuggestionItem
     ? {
         // can be undefined
@@ -67,7 +67,6 @@ export function SuggestionMenuController<
     minQueryLength,
     onItemClick,
     getItems,
-    floatingOptions,
   } = props;
 
   const onItemClickOrDefault = useMemo(() => {
@@ -90,57 +89,73 @@ export function SuggestionMenuController<
     );
   }, [editor, getItems])!;
 
-  const callbacks = {
-    closeMenu: editor.suggestionMenus.closeMenu,
-    clearQuery: editor.suggestionMenus.clearQuery,
-  };
+  const suggestionMenu = useExtension(SuggestionMenuExtension);
 
-  const cb = useCallback(
-    (callback: (state: SuggestionMenuState) => void) => {
-      return editor.suggestionMenus.onUpdate(triggerCharacter, callback);
-    },
-    [editor.suggestionMenus, triggerCharacter],
-  );
+  useEffect(() => {
+    suggestionMenu.addTriggerCharacter(triggerCharacter);
+  }, [
+    props.triggerCharacter,
+    suggestionMenu,
+    suggestionMenu.store,
+    triggerCharacter,
+  ]);
 
-  const state = useUIPluginState(cb);
+  const state = useExtensionState(SuggestionMenuExtension);
+  const referencePos = useExtensionState(SuggestionMenuExtension, {
+    selector: (state) => state?.referencePos || new DOMRect(),
+  });
 
-  const { isMounted, ref, style, getFloatingProps } = useUIElementPositioning(
-    state?.show || false,
-    state?.referencePos || null,
-    2000,
-    {
-      placement: "bottom-start",
-      middleware: [
-        offset(10),
-        // Flips the menu placement to maximize the space available, and prevents
-        // the menu from being cut off by the confines of the screen.
-        flip({
-          mainAxis: true,
-          crossAxis: false,
-        }),
-        shift(),
-        size({
-          apply({ availableHeight, elements }) {
-            Object.assign(elements.floating.style, {
-              maxHeight: `${availableHeight - 10}px`,
-              minHeight: "300px",
-            });
-          },
-        }),
-      ],
-      onOpenChange(open) {
-        if (!open) {
-          editor.suggestionMenus.closeMenu();
-        }
+  const virtualElement = useMemo<VirtualElement>(
+    () => ({
+      getBoundingClientRect: () => referencePos,
+    }),
+    [referencePos],
+  );
+
+  const floatingUIOptions = useMemo<FloatingUIOptions>(
+    () => ({
+      useFloatingOptions: {
+        open: state?.show && state?.triggerCharacter === triggerCharacter,
+        placement: "bottom-start",
+        middleware: [
+          offset(10),
+          // Flips the menu placement to maximize the space available, and prevents
+          // the menu from being cut off by the confines of the screen.
+          flip({
+            mainAxis: true,
+            crossAxis: false,
+          }),
+          shift(),
+          size({
+            apply({ availableHeight, elements }) {
+              Object.assign(elements.floating.style, {
+                maxHeight: `${availableHeight - 10}px`,
+                minHeight: "300px",
+              });
+            },
+          }),
+        ],
+      },
+      elementProps: {
+        // Prevents editor blurring when clicking the scroll bar.
+        onMouseDownCapture: (event) => event.preventDefault(),
+        style: {
+          zIndex: 80,
+        },
       },
-      ...floatingOptions,
-    },
+      ...props.floatingUIOptions,
+    }),
+    [
+      props.floatingUIOptions,
+      state?.show,
+      state?.triggerCharacter,
+      triggerCharacter,
+    ],
   );
 
   if (
-    !isMounted ||
     !state ||
-    (!state?.ignoreQueryLength &&
+    (!state.ignoreQueryLength &&
       minQueryLength &&
       (state.query.startsWith(" ") || state.query.length < minQueryLength))
   ) {
@@ -148,23 +163,19 @@ export function SuggestionMenuController<
   }
 
   return (
-    <div
-      ref={ref}
-      style={style}
-      {...getFloatingProps()}
-      // Prevents editor blurring when clicking the scroll bar.
-      onMouseDown={(e) => e.preventDefault()}
-    >
-      <SuggestionMenuWrapper
-        query={state.query}
-        closeMenu={callbacks.closeMenu}
-        clearQuery={callbacks.clearQuery}
-        getItems={getItemsOrDefault}
-        suggestionMenuComponent={
-          suggestionMenuComponent || SuggestionMenu<ItemType<GetItemsType>>
-        }
-        onItemClick={onItemClickOrDefault}
-      />
-    </div>
+    <GenericPopover reference={virtualElement} {...floatingUIOptions}>
+      {triggerCharacter && (
+        <SuggestionMenuWrapper
+          query={state.query}
+          closeMenu={suggestionMenu.closeMenu}
+          clearQuery={suggestionMenu.clearQuery}
+          getItems={getItemsOrDefault}
+          suggestionMenuComponent={
+            suggestionMenuComponent || SuggestionMenu<ItemType<GetItemsType>>
+          }
+          onItemClick={onItemClickOrDefault}
+        />
+      )}
+    </GenericPopover>
   );
 }
diff --git a/packages/react/src/components/SuggestionMenu/getDefaultReactSlashMenuItems.tsx b/packages/react/src/components/SuggestionMenu/getDefaultReactSlashMenuItems.tsx
index 17bc55a293..0c62e30f23 100644
--- a/packages/react/src/components/SuggestionMenu/getDefaultReactSlashMenuItems.tsx
+++ b/packages/react/src/components/SuggestionMenu/getDefaultReactSlashMenuItems.tsx
@@ -1,10 +1,10 @@
 import {
   BlockNoteEditor,
   BlockSchema,
-  getDefaultSlashMenuItems,
   InlineContentSchema,
   StyleSchema,
 } from "@blocknote/core";
+import { getDefaultSlashMenuItems } from "@blocknote/core/extensions";
 import {
   RiCodeBlock,
   RiEmotionFill,
diff --git a/packages/react/src/components/SuggestionMenu/types.tsx b/packages/react/src/components/SuggestionMenu/types.tsx
index 8367fd98ea..24647406ff 100644
--- a/packages/react/src/components/SuggestionMenu/types.tsx
+++ b/packages/react/src/components/SuggestionMenu/types.tsx
@@ -1,4 +1,4 @@
-import { DefaultSuggestionItem } from "@blocknote/core";
+import { DefaultSuggestionItem } from "@blocknote/core/extensions";
 import { JSX } from "react";
 
 /**
diff --git a/packages/react/src/components/TableHandles/ExtendButton/ExtendButton.tsx b/packages/react/src/components/TableHandles/ExtendButton/ExtendButton.tsx
index 8e8d499d50..03d167002c 100644
--- a/packages/react/src/components/TableHandles/ExtendButton/ExtendButton.tsx
+++ b/packages/react/src/components/TableHandles/ExtendButton/ExtendButton.tsx
@@ -1,13 +1,10 @@
 import {
-  DefaultInlineContentSchema,
-  DefaultStyleSchema,
   EMPTY_CELL_HEIGHT,
   EMPTY_CELL_WIDTH,
-  InlineContentSchema,
   mergeCSSClasses,
   PartialTableContent,
-  StyleSchema,
 } from "@blocknote/core";
+import { TableHandles } from "@blocknote/core/extensions";
 import {
   MouseEvent as ReactMouseEvent,
   ReactNode,
@@ -19,6 +16,11 @@ import {
 import { RiAddFill } from "react-icons/ri";
 
 import { useComponentsContext } from "../../../editor/ComponentsContext.js";
+import { useBlockNoteEditor } from "../../../hooks/useBlockNoteEditor.js";
+import {
+  useExtension,
+  useExtensionState,
+} from "../../../hooks/useExtension.js";
 import { ExtendButtonProps } from "./ExtendButtonProps.js";
 
 // Rounds a number up or down, depending on whether we're close (as defined by
@@ -36,22 +38,26 @@ const marginRound = (num: number, margin = 0.3) => {
   }
 };
 
-export const ExtendButton = <
-  I extends InlineContentSchema = DefaultInlineContentSchema,
-  S extends StyleSchema = DefaultStyleSchema,
->(
-  props: ExtendButtonProps<I, S> & { children?: ReactNode },
+export const ExtendButton = (
+  props: ExtendButtonProps & { children?: ReactNode },
 ) => {
   const Components = useComponentsContext()!;
 
+  const editor = useBlockNoteEditor<any, any, any>();
+
+  const tableHandles = useExtension(TableHandles);
+  const block = useExtensionState(TableHandles, {
+    selector: (state) => state?.block,
+  });
+
   // needs to be a ref because it's used immediately in the onClick handler
   // (state would be async and only have effect after the next render
   const movedMouse = useRef(false);
 
   const [editingState, setEditingState] = useState<
     | {
-        originalContent: PartialTableContent<I, S>;
-        originalCroppedContent: PartialTableContent<I, S>;
+        originalContent: PartialTableContent<any, any>;
+        originalCroppedContent: PartialTableContent<any, any>;
         startPos: number;
       }
     | undefined
@@ -60,12 +66,18 @@ export const ExtendButton = <
   // Lets the user start extending columns/rows by moving the mouse.
   const mouseDownHandler = useCallback(
     (event: ReactMouseEvent) => {
-      props.onMouseDown();
+      tableHandles.freezeHandles();
+      props.hideOtherElements(true);
+
+      if (!block) {
+        return;
+      }
+
       setEditingState({
-        originalContent: props.block.content,
+        originalContent: block.content as any,
         originalCroppedContent: {
-          rows: props.editor.tableHandles!.cropEmptyRowsOrColumns(
-            props.block,
+          rows: tableHandles.cropEmptyRowsOrColumns(
+            block,
             props.orientation === "addOrRemoveColumns" ? "columns" : "rows",
           ),
         } as PartialTableContent<any, any>,
@@ -79,36 +91,33 @@ export const ExtendButton = <
       // preventdefault, otherwise text in the table might be selected
       event.preventDefault();
     },
-    [props],
+    [block, props, tableHandles],
   );
 
   const onClickHandler = useCallback(() => {
-    if (movedMouse.current) {
+    if (!block || movedMouse.current) {
       return;
     }
-    props.editor.updateBlock(props.block, {
+
+    editor.updateBlock(block, {
       type: "table",
       content: {
-        ...props.block.content,
+        ...block.content,
         rows:
           props.orientation === "addOrRemoveColumns"
-            ? props.editor.tableHandles!.addRowsOrColumns(
-                props.block,
-                "columns",
-                1,
-              )
-            : props.editor.tableHandles!.addRowsOrColumns(
-                props.block,
-                "rows",
-                1,
-              ),
+            ? tableHandles.addRowsOrColumns(block, "columns", 1)
+            : tableHandles.addRowsOrColumns(block, "rows", 1),
       } as any,
     });
-  }, [props.block, props.orientation, props.editor]);
+  }, [block, editor, props.orientation, tableHandles]);
 
   // Extends columns/rows on when moving the mouse.
   useEffect(() => {
     const callback = (event: MouseEvent) => {
+      if (!block) {
+        return;
+      }
+
       if (!editingState) {
         throw new Error("editingState is undefined");
       }
@@ -132,8 +141,8 @@ export const ExtendButton = <
 
       const currentNumCells =
         props.orientation === "addOrRemoveColumns"
-          ? props.block.content.rows[0].cells.length
-          : props.block.content.rows.length;
+          ? block.content.rows[0].cells.length
+          : block.content.rows.length;
 
       const newNumCells =
         numOriginalCells +
@@ -150,13 +159,13 @@ export const ExtendButton = <
         newNumCells > 0 &&
         newNumCells !== currentNumCells
       ) {
-        props.editor.updateBlock(props.block, {
+        editor.updateBlock(block, {
           type: "table",
           content: {
-            ...props.block.content,
+            ...block.content,
             rows:
               props.orientation === "addOrRemoveColumns"
-                ? props.editor.tableHandles!.addRowsOrColumns(
+                ? tableHandles.addRowsOrColumns(
                     {
                       type: "table",
                       content: editingState.originalCroppedContent,
@@ -164,7 +173,7 @@ export const ExtendButton = <
                     "columns",
                     newNumCells - numCroppedCells,
                   )
-                : props.editor.tableHandles!.addRowsOrColumns(
+                : tableHandles.addRowsOrColumns(
                     {
                       type: "table",
                       content: editingState.originalCroppedContent,
@@ -177,8 +186,8 @@ export const ExtendButton = <
 
         // Edge case for updating block content as `updateBlock` causes the
         // selection to move into the next block, so we have to set it back.
-        if (props.block.content) {
-          props.editor.setTextCursorPosition(props.block);
+        if (block.content) {
+          editor.setTextCursorPosition(block);
         }
       }
     };
@@ -186,20 +195,20 @@ export const ExtendButton = <
     if (editingState) {
       window.addEventListener("mousemove", callback);
     }
+
     return () => {
       window.removeEventListener("mousemove", callback);
     };
-  }, [editingState, props.block, props.editor, props.orientation]);
+  }, [block, editingState, editor, props.orientation, tableHandles]);
 
   // Stops mouse movements from extending columns/rows when the mouse is
   // released. Also extends columns/rows by 1 if the mouse wasn't moved enough
   // to add any, imitating a click.
   useEffect(() => {
-    const onMouseUp = props.onMouseUp;
-
     const callback = () => {
+      props.hideOtherElements(false);
+      tableHandles.unfreezeHandles();
       setEditingState(undefined);
-      onMouseUp();
     };
 
     if (editingState) {
@@ -209,9 +218,9 @@ export const ExtendButton = <
     return () => {
       window.removeEventListener("mouseup", callback);
     };
-  }, [editingState, props.onMouseUp]);
+  }, [editingState, props, tableHandles]);
 
-  if (!props.editor.isEditable) {
+  if (!editor.isEditable) {
     return null;
   }
 
diff --git a/packages/react/src/components/TableHandles/ExtendButton/ExtendButtonProps.ts b/packages/react/src/components/TableHandles/ExtendButton/ExtendButtonProps.ts
index ef3104de27..face15b7e7 100644
--- a/packages/react/src/components/TableHandles/ExtendButton/ExtendButtonProps.ts
+++ b/packages/react/src/components/TableHandles/ExtendButton/ExtendButtonProps.ts
@@ -1,25 +1,4 @@
-import {
-  BlockNoteEditor,
-  DefaultBlockSchema,
-  DefaultInlineContentSchema,
-  DefaultStyleSchema,
-  InlineContentSchema,
-  StyleSchema,
-  TableHandlesState,
-} from "@blocknote/core";
-
-export type ExtendButtonProps<
-  I extends InlineContentSchema = DefaultInlineContentSchema,
-  S extends StyleSchema = DefaultStyleSchema,
-> = {
-  editor: BlockNoteEditor<
-    {
-      table: DefaultBlockSchema["table"];
-    },
-    I,
-    S
-  >;
-  onMouseDown: () => void;
-  onMouseUp: () => void;
+export type ExtendButtonProps = {
   orientation: "addOrRemoveRows" | "addOrRemoveColumns";
-} & Pick<TableHandlesState<I, S>, "block">;
+  hideOtherElements: (hide: boolean) => void;
+};
diff --git a/packages/react/src/components/TableHandles/TableCellButton.tsx b/packages/react/src/components/TableHandles/TableCellButton.tsx
index 3a16568eb3..74add39bf6 100644
--- a/packages/react/src/components/TableHandles/TableCellButton.tsx
+++ b/packages/react/src/components/TableHandles/TableCellButton.tsx
@@ -1,14 +1,10 @@
-import {
-  DefaultInlineContentSchema,
-  DefaultStyleSchema,
-  InlineContentSchema,
-  StyleSchema,
-} from "@blocknote/core";
+import { TableHandles } from "@blocknote/core/extensions";
 import { ReactNode } from "react";
-
-import { createPortal } from "react-dom";
 import { MdArrowDropDown } from "react-icons/md";
+
 import { useComponentsContext } from "../../editor/ComponentsContext.js";
+import { useBlockNoteEditor } from "../../hooks/useBlockNoteEditor.js";
+import { useExtension } from "../../hooks/useExtension.js";
 import { TableCellButtonProps } from "./TableCellButtonProps.js";
 import { TableCellMenu } from "./TableCellMenu/TableCellMenu.js";
 
@@ -16,20 +12,21 @@ import { TableCellMenu } from "./TableCellMenu/TableCellMenu.js";
  * By default, the TableCellHandle component will render with the default icon.
  * However, you can override the icon to render by passing children.
  */
-export const TableCellButton = <
-  I extends InlineContentSchema = DefaultInlineContentSchema,
-  S extends StyleSchema = DefaultStyleSchema,
->(
-  props: TableCellButtonProps<I, S> & { children?: ReactNode },
+export const TableCellButton = (
+  props: TableCellButtonProps & { children?: ReactNode },
 ) => {
   const Components = useComponentsContext()!;
 
+  const editor = useBlockNoteEditor<any, any, any>();
+
+  const tableHandles = useExtension(TableHandles);
+
   const Component = props.tableCellMenu || TableCellMenu;
 
   if (
-    !props.editor.settings.tables.splitCells &&
-    !props.editor.settings.tables.cellBackgroundColor &&
-    !props.editor.settings.tables.cellTextColor
+    !editor.settings.tables.splitCells &&
+    !editor.settings.tables.cellBackgroundColor &&
+    !editor.settings.tables.cellTextColor
   ) {
     // Hide the button altogether if all table cell settings are disabled
     return null;
@@ -39,10 +36,12 @@ export const TableCellButton = <
     <Components.Generic.Menu.Root
       onOpenChange={(open: boolean) => {
         if (open) {
-          props.freezeHandles();
+          tableHandles.freezeHandles();
+          props.hideOtherElements(true);
         } else {
-          props.unfreezeHandles();
-          props.editor.focus();
+          tableHandles.unfreezeHandles();
+          props.hideOtherElements(false);
+          editor.focus();
         }
       }}
       position={"right"}
@@ -54,15 +53,7 @@ export const TableCellButton = <
           )}
         </Components.Generic.Menu.Button>
       </Components.Generic.Menu.Trigger>
-      {/* the menu can extend outside of the table, so we use a portal to prevent clipping */}
-      {createPortal(
-        <Component
-          block={props.block as any}
-          rowIndex={props.rowIndex}
-          colIndex={props.colIndex}
-        />,
-        props.menuContainer,
-      )}
+      <Component />
     </Components.Generic.Menu.Root>
   );
 };
diff --git a/packages/react/src/components/TableHandles/TableCellButtonProps.ts b/packages/react/src/components/TableHandles/TableCellButtonProps.ts
index 1f9a307e2a..21dd57425f 100644
--- a/packages/react/src/components/TableHandles/TableCellButtonProps.ts
+++ b/packages/react/src/components/TableHandles/TableCellButtonProps.ts
@@ -1,41 +1,6 @@
-import {
-  BlockNoteEditor,
-  DefaultBlockSchema,
-  DefaultInlineContentSchema,
-  DefaultStyleSchema,
-  InlineContentSchema,
-  StyleSchema,
-  TableHandlesState,
-} from "@blocknote/core";
-import { TableCellMenuProps } from "./TableCellMenu/TableCellMenuProps.js";
 import { FC } from "react";
 
-export type TableCellButtonProps<
-  I extends InlineContentSchema = DefaultInlineContentSchema,
-  S extends StyleSchema = DefaultStyleSchema,
-> = {
-  editor: BlockNoteEditor<
-    {
-      table: DefaultBlockSchema["table"];
-    },
-    I,
-    S
-  >;
-  rowIndex: number;
-  colIndex: number;
-  menuContainer: HTMLDivElement;
-  tableCellMenu?: FC<TableCellMenuProps<I, S>>;
-} & Pick<TableHandlesState<I, S>, "block"> &
-  Pick<
-    Exclude<
-      BlockNoteEditor<
-        {
-          table: DefaultBlockSchema["table"];
-        },
-        I,
-        S
-      >["tableHandles"],
-      undefined
-    >,
-    "freezeHandles" | "unfreezeHandles"
-  >;
+export type TableCellButtonProps = {
+  hideOtherElements: (hide: boolean) => void;
+  tableCellMenu?: FC;
+};
diff --git a/packages/react/src/components/TableHandles/TableCellMenu/DefaultButtons/ColorPicker.tsx b/packages/react/src/components/TableHandles/TableCellMenu/DefaultButtons/ColorPicker.tsx
index eddbc6b5c7..ee423551a6 100644
--- a/packages/react/src/components/TableHandles/TableCellMenu/DefaultButtons/ColorPicker.tsx
+++ b/packages/react/src/components/TableHandles/TableCellMenu/DefaultButtons/ColorPicker.tsx
@@ -1,38 +1,38 @@
-import {
-  DefaultBlockSchema,
-  DefaultInlineContentSchema,
-  DefaultStyleSchema,
-  InlineContentSchema,
-  isTableCell,
-  mapTableCell,
-  StyleSchema,
-} from "@blocknote/core";
+import { isTableCell, mapTableCell } from "@blocknote/core";
+import { TableHandles } from "@blocknote/core/extensions";
+import { ReactNode } from "react";
 
 import { useComponentsContext } from "../../../../editor/ComponentsContext.js";
 import { useBlockNoteEditor } from "../../../../hooks/useBlockNoteEditor.js";
+import { useExtensionState } from "../../../../hooks/useExtension.js";
 import { useDictionary } from "../../../../i18n/dictionary.js";
 import { ColorPicker } from "../../../ColorPicker/ColorPicker.js";
-import { TableCellMenuProps } from "../TableCellMenuProps.js";
-import { ReactNode } from "react";
 
-export const ColorPickerButton = <
-  I extends InlineContentSchema = DefaultInlineContentSchema,
-  S extends StyleSchema = DefaultStyleSchema,
->(
-  props: TableCellMenuProps<I, S> & {
-    children?: ReactNode;
-  },
-) => {
+export const ColorPickerButton = (props: { children?: ReactNode }) => {
   const Components = useComponentsContext()!;
   const dict = useDictionary();
-  const editor = useBlockNoteEditor<
-    { table: DefaultBlockSchema["table"] },
-    I,
-    S
-  >();
+  const editor = useBlockNoteEditor<any, any, any>();
+
+  const block = useExtensionState(TableHandles, {
+    selector: (state) => state?.block,
+  });
+  const colIndex = useExtensionState(TableHandles, {
+    selector: (state) => state?.colIndex,
+  });
+  const rowIndex = useExtensionState(TableHandles, {
+    selector: (state) => state?.rowIndex,
+  });
 
   const updateColor = (color: string, type: "text" | "background") => {
-    const newTable = props.block.content.rows.map((row) => {
+    if (
+      block === undefined ||
+      colIndex === undefined ||
+      rowIndex === undefined
+    ) {
+      return;
+    }
+
+    const newTable = block.content.rows.map((row) => {
       return {
         ...row,
         cells: row.cells.map((cell) => mapTableCell(cell)),
@@ -40,27 +40,29 @@ export const ColorPickerButton = <
     });
 
     if (type === "text") {
-      newTable[props.rowIndex].cells[props.colIndex].props.textColor = color;
+      newTable[rowIndex].cells[colIndex].props.textColor = color;
     } else {
-      newTable[props.rowIndex].cells[props.colIndex].props.backgroundColor =
-        color;
+      newTable[rowIndex].cells[colIndex].props.backgroundColor = color;
     }
 
-    editor.updateBlock(props.block, {
+    editor.updateBlock(block, {
       type: "table",
       content: {
-        ...props.block.content,
+        ...block.content,
         rows: newTable,
-      },
+      } as any,
     });
 
     // Have to reset text cursor position to the block as `updateBlock`
     // moves the existing selection out of the block.
-    editor.setTextCursorPosition(props.block);
+    editor.setTextCursorPosition(block);
   };
 
-  const currentCell =
-    props.block.content.rows[props.rowIndex]?.cells?.[props.colIndex];
+  if (block === undefined || colIndex === undefined || rowIndex === undefined) {
+    return null;
+  }
+
+  const currentCell = block.content.rows[rowIndex]?.cells?.[colIndex];
 
   if (
     !currentCell ||
diff --git a/packages/react/src/components/TableHandles/TableCellMenu/DefaultButtons/SplitButton.tsx b/packages/react/src/components/TableHandles/TableCellMenu/DefaultButtons/SplitButton.tsx
index 4f04f2cea6..397c299c20 100644
--- a/packages/react/src/components/TableHandles/TableCellMenu/DefaultButtons/SplitButton.tsx
+++ b/packages/react/src/components/TableHandles/TableCellMenu/DefaultButtons/SplitButton.tsx
@@ -1,35 +1,35 @@
-import {
-  DefaultBlockSchema,
-  DefaultInlineContentSchema,
-  DefaultStyleSchema,
-  getColspan,
-  getRowspan,
-  InlineContentSchema,
-  isTableCell,
-  StyleSchema,
-} from "@blocknote/core";
+import { getColspan, getRowspan, isTableCell } from "@blocknote/core";
+import { TableHandles } from "@blocknote/core/extensions";
 
 import { useComponentsContext } from "../../../../editor/ComponentsContext.js";
 import { useBlockNoteEditor } from "../../../../hooks/useBlockNoteEditor.js";
 import { useDictionary } from "../../../../i18n/dictionary.js";
-import { TableCellMenuProps } from "../TableCellMenuProps.js";
+import {
+  useExtension,
+  useExtensionState,
+} from "../../../../hooks/useExtension.js";
 
-export const SplitButton = <
-  I extends InlineContentSchema = DefaultInlineContentSchema,
-  S extends StyleSchema = DefaultStyleSchema,
->(
-  props: TableCellMenuProps<I, S>,
-) => {
+export const SplitButton = () => {
   const Components = useComponentsContext()!;
   const dict = useDictionary();
-  const editor = useBlockNoteEditor<
-    { table: DefaultBlockSchema["table"] },
-    I,
-    S
-  >();
+  const editor = useBlockNoteEditor<any, any, any>();
+
+  const tableHandles = useExtension(TableHandles);
+  const block = useExtensionState(TableHandles, {
+    selector: (state) => state?.block,
+  });
+  const colIndex = useExtensionState(TableHandles, {
+    selector: (state) => state?.colIndex,
+  });
+  const rowIndex = useExtensionState(TableHandles, {
+    selector: (state) => state?.rowIndex,
+  });
+
+  if (block === undefined || colIndex === undefined || rowIndex === undefined) {
+    return null;
+  }
 
-  const currentCell =
-    props.block.content.rows[props.rowIndex]?.cells?.[props.colIndex];
+  const currentCell = block.content.rows[rowIndex]?.cells?.[colIndex];
 
   if (
     !currentCell ||
@@ -43,9 +43,9 @@ export const SplitButton = <
   return (
     <Components.Generic.Menu.Item
       onClick={() => {
-        editor.tableHandles?.splitCell({
-          row: props.rowIndex,
-          col: props.colIndex,
+        tableHandles.splitCell({
+          row: rowIndex,
+          col: colIndex,
         });
       }}
     >
diff --git a/packages/react/src/components/TableHandles/TableCellMenu/TableCellMenu.tsx b/packages/react/src/components/TableHandles/TableCellMenu/TableCellMenu.tsx
index 84eebba026..f8722ad6d1 100644
--- a/packages/react/src/components/TableHandles/TableCellMenu/TableCellMenu.tsx
+++ b/packages/react/src/components/TableHandles/TableCellMenu/TableCellMenu.tsx
@@ -1,22 +1,10 @@
-import {
-  DefaultInlineContentSchema,
-  DefaultStyleSchema,
-  InlineContentSchema,
-  StyleSchema,
-} from "@blocknote/core";
 import { ReactNode } from "react";
 
 import { useComponentsContext } from "../../../editor/ComponentsContext.js";
 import { ColorPickerButton } from "./DefaultButtons/ColorPicker.js";
 import { SplitButton } from "./DefaultButtons/SplitButton.js";
-import { TableCellMenuProps } from "./TableCellMenuProps.js";
 
-export const TableCellMenu = <
-  I extends InlineContentSchema = DefaultInlineContentSchema,
-  S extends StyleSchema = DefaultStyleSchema,
->(
-  props: TableCellMenuProps<I, S> & { children?: ReactNode },
-) => {
+export const TableCellMenu = (props: { children?: ReactNode }) => {
   const Components = useComponentsContext()!;
 
   return (
@@ -25,16 +13,8 @@ export const TableCellMenu = <
     >
       {props.children || (
         <>
-          <SplitButton
-            block={props.block}
-            rowIndex={props.rowIndex}
-            colIndex={props.colIndex}
-          />
-          <ColorPickerButton
-            block={props.block}
-            rowIndex={props.rowIndex}
-            colIndex={props.colIndex}
-          />
+          <SplitButton />
+          <ColorPickerButton />
         </>
       )}
     </Components.Generic.Menu.Dropdown>
diff --git a/packages/react/src/components/TableHandles/TableCellMenu/TableCellMenuProps.ts b/packages/react/src/components/TableHandles/TableCellMenu/TableCellMenuProps.ts
deleted file mode 100644
index b2af7056d4..0000000000
--- a/packages/react/src/components/TableHandles/TableCellMenu/TableCellMenuProps.ts
+++ /dev/null
@@ -1,17 +0,0 @@
-import {
-  DefaultBlockSchema,
-  DefaultInlineContentSchema,
-  DefaultStyleSchema,
-  InlineContentSchema,
-  SpecificBlock,
-  StyleSchema,
-} from "@blocknote/core";
-
-export type TableCellMenuProps<
-  I extends InlineContentSchema = DefaultInlineContentSchema,
-  S extends StyleSchema = DefaultStyleSchema,
-> = {
-  block: SpecificBlock<{ table: DefaultBlockSchema["table"] }, "table", I, S>;
-  rowIndex: number;
-  colIndex: number;
-};
diff --git a/packages/react/src/components/TableHandles/TableHandle.tsx b/packages/react/src/components/TableHandles/TableHandle.tsx
index a4611eba82..9f7b62dc76 100644
--- a/packages/react/src/components/TableHandles/TableHandle.tsx
+++ b/packages/react/src/components/TableHandles/TableHandle.tsx
@@ -1,17 +1,11 @@
-import {
-  DefaultInlineContentSchema,
-  DefaultStyleSchema,
-  getColspan,
-  getRowspan,
-  InlineContentSchema,
-  mergeCSSClasses,
-  StyleSchema,
-} from "@blocknote/core";
+import { getColspan, getRowspan, mergeCSSClasses } from "@blocknote/core";
+import { TableHandles } from "@blocknote/core/extensions";
 import { ReactNode, useMemo, useState } from "react";
 
-import { createPortal } from "react-dom";
 import { MdDragIndicator } from "react-icons/md";
 import { useComponentsContext } from "../../editor/ComponentsContext.js";
+import { useBlockNoteEditor } from "../../hooks/useBlockNoteEditor.js";
+import { useExtension, useExtensionState } from "../../hooks/useExtension.js";
 import { TableHandleMenu } from "./TableHandleMenu/TableHandleMenu.js";
 import { TableHandleProps } from "./TableHandleProps.js";
 
@@ -19,45 +13,56 @@ import { TableHandleProps } from "./TableHandleProps.js";
  * By default, the TableHandle component will render with the default icon.
  * However, you can override the icon to render by passing children.
  */
-export const TableHandle = <
-  I extends InlineContentSchema = DefaultInlineContentSchema,
-  S extends StyleSchema = DefaultStyleSchema,
->(
-  props: TableHandleProps<I, S> & { children?: ReactNode },
+export const TableHandle = (
+  props: TableHandleProps & {
+    children?: ReactNode;
+  },
 ) => {
+  const editor = useBlockNoteEditor<any, any, any>();
   const Components = useComponentsContext()!;
 
   const [isDragging, setIsDragging] = useState(false);
 
   const Component = props.tableHandleMenu || TableHandleMenu;
 
+  const tableHandles = useExtension(TableHandles);
+  const state = useExtensionState(TableHandles);
+
   const isDraggable = useMemo(() => {
-    const tableHandles = props.editor.tableHandles;
-    if (!tableHandles || !props.block || props.block.type !== "table") {
+    if (
+      !tableHandles ||
+      !state ||
+      !state.block ||
+      state.block.type !== "table"
+    ) {
       return false;
     }
 
     if (props.orientation === "column") {
       return tableHandles
-        .getCellsAtColumnHandle(props.block, props.index)
+        .getCellsAtColumnHandle(state.block, state.colIndex!)
         .every(({ cell }) => getColspan(cell) === 1);
     }
 
     return tableHandles
-      .getCellsAtRowHandle(props.block, props.index)
+      .getCellsAtRowHandle(state.block, state.rowIndex!)
       .every(({ cell }) => getRowspan(cell) === 1);
-  }, [props.block, props.editor.tableHandles, props.index, props.orientation]);
+  }, [props.orientation, state, tableHandles]);
+
+  if (!state) {
+    return null;
+  }
 
   return (
     <Components.Generic.Menu.Root
       onOpenChange={(open: boolean) => {
         if (open) {
-          props.freezeHandles();
-          props.hideOtherSide();
+          tableHandles.freezeHandles();
+          props.hideOtherElements(true);
         } else {
-          props.unfreezeHandles();
-          props.showOtherSide();
-          props.editor.focus();
+          tableHandles.unfreezeHandles();
+          props.hideOtherElements(false);
+          editor.focus();
         }
       }}
       position={"right"}
@@ -72,10 +77,16 @@ export const TableHandle = <
           draggable={isDraggable}
           onDragStart={(e) => {
             setIsDragging(true);
-            props.dragStart(e);
+            props.hideOtherElements(true);
+            if (props.orientation === "column") {
+              tableHandles.colDragStart(e);
+            } else {
+              tableHandles.rowDragStart(e);
+            }
           }}
           onDragEnd={() => {
-            props.dragEnd();
+            tableHandles.dragEnd();
+            props.hideOtherElements(false);
             setIsDragging(false);
           }}
           style={
@@ -89,15 +100,7 @@ export const TableHandle = <
           )}
         </Components.TableHandle.Root>
       </Components.Generic.Menu.Trigger>
-      {/* the menu can extend outside of the table, so we use a portal to prevent clipping */}
-      {createPortal(
-        <Component
-          orientation={props.orientation}
-          block={props.block as any}
-          index={props.index}
-        />,
-        props.menuContainer,
-      )}
+      <Component orientation={props.orientation} />
     </Components.Generic.Menu.Root>
   );
 };
diff --git a/packages/react/src/components/TableHandles/TableHandleMenu/DefaultButtons/AddButton.tsx b/packages/react/src/components/TableHandles/TableHandleMenu/DefaultButtons/AddButton.tsx
index b4e0b56e7b..23bee820d8 100644
--- a/packages/react/src/components/TableHandles/TableHandleMenu/DefaultButtons/AddButton.tsx
+++ b/packages/react/src/components/TableHandles/TableHandleMenu/DefaultButtons/AddButton.tsx
@@ -1,38 +1,27 @@
-import {
-  DefaultBlockSchema,
-  DefaultInlineContentSchema,
-  DefaultStyleSchema,
-  InlineContentSchema,
-  StyleSchema,
-} from "@blocknote/core";
+import { TableHandles } from "@blocknote/core/extensions";
 
 import { useComponentsContext } from "../../../../editor/ComponentsContext.js";
-import { useBlockNoteEditor } from "../../../../hooks/useBlockNoteEditor.js";
 import { useDictionary } from "../../../../i18n/dictionary.js";
-import { TableHandleMenuProps } from "../TableHandleMenuProps.js";
+import {
+  useExtension,
+  useExtensionState,
+} from "../../../../hooks/useExtension.js";
 
-export const AddButton = <
-  I extends InlineContentSchema = DefaultInlineContentSchema,
-  S extends StyleSchema = DefaultStyleSchema,
->(
-  props: TableHandleMenuProps<I, S> &
-    (
-      | { orientation: "row"; side: "above" | "below" }
-      | { orientation: "column"; side: "left" | "right" }
-    ),
+export const AddButton = (
+  props:
+    | { orientation: "row"; side: "above" | "below" }
+    | { orientation: "column"; side: "left" | "right" },
 ) => {
   const Components = useComponentsContext()!;
   const dict = useDictionary();
 
-  const editor = useBlockNoteEditor<
-    { table: DefaultBlockSchema["table"] },
-    I,
-    S
-  >();
-
-  const tableHandles = editor.tableHandles;
+  const tableHandles = useExtension(TableHandles);
+  const index = useExtensionState(TableHandles, {
+    selector: (state) =>
+      props.orientation === "column" ? state?.colIndex : state?.rowIndex,
+  });
 
-  if (!tableHandles) {
+  if (tableHandles === undefined || index === undefined) {
     return null;
   }
 
@@ -40,7 +29,7 @@ export const AddButton = <
     <Components.Generic.Menu.Item
       onClick={() => {
         tableHandles.addRowOrColumn(
-          props.index,
+          index,
           props.orientation === "row"
             ? { orientation: "row", side: props.side }
             : { orientation: "column", side: props.side },
diff --git a/packages/react/src/components/TableHandles/TableHandleMenu/DefaultButtons/ColorPicker.tsx b/packages/react/src/components/TableHandles/TableHandleMenu/DefaultButtons/ColorPicker.tsx
index b93d653fa1..0d2b4fe654 100644
--- a/packages/react/src/components/TableHandles/TableHandleMenu/DefaultButtons/ColorPicker.tsx
+++ b/packages/react/src/components/TableHandles/TableHandleMenu/DefaultButtons/ColorPicker.tsx
@@ -7,22 +7,25 @@ import {
   mapTableCell,
   StyleSchema,
 } from "@blocknote/core";
+import { TableHandles } from "@blocknote/core/extensions";
 
 import { useComponentsContext } from "../../../../editor/ComponentsContext.js";
 import { useBlockNoteEditor } from "../../../../hooks/useBlockNoteEditor.js";
 import { useDictionary } from "../../../../i18n/dictionary.js";
 import { ColorPicker } from "../../../ColorPicker/ColorPicker.js";
-import { TableHandleMenuProps } from "../TableHandleMenuProps.js";
 import { ReactNode, useMemo } from "react";
+import {
+  useExtension,
+  useExtensionState,
+} from "../../../../hooks/useExtension.js";
 
 export const ColorPickerButton = <
   I extends InlineContentSchema = DefaultInlineContentSchema,
   S extends StyleSchema = DefaultStyleSchema,
->(
-  props: TableHandleMenuProps<I, S> & {
-    children?: ReactNode;
-  },
-) => {
+>(props: {
+  orientation: "row" | "column";
+  children?: ReactNode;
+}) => {
   const Components = useComponentsContext()!;
   const dict = useDictionary();
   const editor = useBlockNoteEditor<
@@ -30,22 +33,38 @@ export const ColorPickerButton = <
     I,
     S
   >();
-  const tableHandles = editor.tableHandles;
+
+  const tableHandles = useExtension(TableHandles);
+  const block = useExtensionState(TableHandles, {
+    selector: (state) => state?.block,
+  });
+  const index = useExtensionState(TableHandles, {
+    selector: (state) =>
+      props.orientation === "column" ? state?.colIndex : state?.rowIndex,
+  });
 
   const currentCells = useMemo(() => {
-    if (!tableHandles || !props.block) {
+    if (
+      tableHandles === undefined ||
+      block === undefined ||
+      index === undefined
+    ) {
       return [];
     }
 
     if (props.orientation === "row") {
-      return tableHandles.getCellsAtRowHandle(props.block, props.index);
+      return tableHandles.getCellsAtRowHandle(block, index);
     }
 
-    return tableHandles.getCellsAtColumnHandle(props.block, props.index);
-  }, [props.block, props.index, props.orientation, tableHandles]);
+    return tableHandles.getCellsAtColumnHandle(block, index);
+  }, [block, index, props.orientation, tableHandles]);
 
   const updateColor = (color: string, type: "text" | "background") => {
-    const newTable = props.block.content.rows.map((row) => {
+    if (block === undefined) {
+      return;
+    }
+
+    const newTable = block.content.rows.map((row) => {
       return {
         ...row,
         cells: row.cells.map((cell) => mapTableCell(cell)),
@@ -60,17 +79,17 @@ export const ColorPickerButton = <
       }
     });
 
-    editor.updateBlock(props.block, {
+    editor.updateBlock(block, {
       type: "table",
       content: {
-        ...props.block.content,
+        ...block.content,
         rows: newTable,
-      },
+      } as any,
     });
 
     // Have to reset text cursor position to the block as `updateBlock`
     // moves the existing selection out of the block.
-    editor.setTextCursorPosition(props.block);
+    editor.setTextCursorPosition(block);
   };
 
   if (
diff --git a/packages/react/src/components/TableHandles/TableHandleMenu/DefaultButtons/DeleteButton.tsx b/packages/react/src/components/TableHandles/TableHandleMenu/DefaultButtons/DeleteButton.tsx
index a51474f82e..f16d373362 100644
--- a/packages/react/src/components/TableHandles/TableHandleMenu/DefaultButtons/DeleteButton.tsx
+++ b/packages/react/src/components/TableHandles/TableHandleMenu/DefaultButtons/DeleteButton.tsx
@@ -1,41 +1,29 @@
-import {
-  DefaultBlockSchema,
-  DefaultInlineContentSchema,
-  DefaultStyleSchema,
-  InlineContentSchema,
-  StyleSchema,
-} from "@blocknote/core";
+import { TableHandles } from "@blocknote/core/extensions";
 
 import { useComponentsContext } from "../../../../editor/ComponentsContext.js";
-import { useBlockNoteEditor } from "../../../../hooks/useBlockNoteEditor.js";
 import { useDictionary } from "../../../../i18n/dictionary.js";
-import { TableHandleMenuProps } from "../TableHandleMenuProps.js";
+import {
+  useExtension,
+  useExtensionState,
+} from "../../../../hooks/useExtension.js";
 
-export const DeleteButton = <
-  I extends InlineContentSchema = DefaultInlineContentSchema,
-  S extends StyleSchema = DefaultStyleSchema,
->(
-  props: TableHandleMenuProps<I, S> & { orientation: "row" | "column" },
-) => {
+export const DeleteButton = (props: { orientation: "row" | "column" }) => {
   const Components = useComponentsContext()!;
   const dict = useDictionary();
-  const editor = useBlockNoteEditor<
-    { table: DefaultBlockSchema["table"] },
-    I,
-    S
-  >();
 
-  const tableHandles = editor.tableHandles;
+  const tableHandles = useExtension(TableHandles);
+  const index = useExtensionState(TableHandles, {
+    selector: (state) =>
+      props.orientation === "column" ? state?.colIndex : state?.rowIndex,
+  });
 
-  if (!tableHandles) {
+  if (tableHandles === undefined || index === undefined) {
     return null;
   }
 
   return (
     <Components.Generic.Menu.Item
-      onClick={() => {
-        tableHandles.removeRowOrColumn(props.index, props.orientation);
-      }}
+      onClick={() => tableHandles.removeRowOrColumn(index, props.orientation)}
     >
       {props.orientation === "row"
         ? dict.table_handle.delete_row_menuitem
diff --git a/packages/react/src/components/TableHandles/TableHandleMenu/DefaultButtons/TableHeaderButton.tsx b/packages/react/src/components/TableHandles/TableHandleMenu/DefaultButtons/TableHeaderButton.tsx
index 7dbb8c4e7c..a72786a7f4 100644
--- a/packages/react/src/components/TableHandles/TableHandleMenu/DefaultButtons/TableHeaderButton.tsx
+++ b/packages/react/src/components/TableHandles/TableHandleMenu/DefaultButtons/TableHeaderButton.tsx
@@ -5,18 +5,22 @@ import {
   InlineContentSchema,
   StyleSchema,
 } from "@blocknote/core";
+import { TableHandles } from "@blocknote/core/extensions";
 
 import { useComponentsContext } from "../../../../editor/ComponentsContext.js";
 import { useBlockNoteEditor } from "../../../../hooks/useBlockNoteEditor.js";
-import { TableHandleMenuProps } from "../TableHandleMenuProps.js";
+import {
+  useExtension,
+  useExtensionState,
+} from "../../../../hooks/useExtension.js";
 import { useDictionary } from "../../../../i18n/dictionary.js";
 
 export const TableHeaderRowButton = <
   I extends InlineContentSchema = DefaultInlineContentSchema,
   S extends StyleSchema = DefaultStyleSchema,
->(
-  props: TableHandleMenuProps<I, S> & { orientation: "row" | "column" },
-) => {
+>(props: {
+  orientation: "row" | "column";
+}) => {
   const Components = useComponentsContext()!;
   const dict = useDictionary();
 
@@ -25,11 +29,19 @@ export const TableHeaderRowButton = <
     I,
     S
   >();
-  const tableHandles = editor.tableHandles;
+  const tableHandles = useExtension(TableHandles);
+  const block = useExtensionState(TableHandles, {
+    selector: (state) => state?.block,
+  });
+  const index = useExtensionState(TableHandles, {
+    selector: (state) =>
+      props.orientation === "column" ? state?.colIndex : state?.rowIndex,
+  });
 
   if (
-    !tableHandles ||
-    props.index !== 0 ||
+    tableHandles === undefined ||
+    block === undefined ||
+    index !== 0 ||
     props.orientation !== "row" ||
     !editor.settings.tables.headers
   ) {
@@ -37,24 +49,19 @@ export const TableHeaderRowButton = <
   }
 
   // We only support 1 header row for now
-  const isHeaderRow = Boolean(props.block.content.headerRows);
+  const isHeaderRow = Boolean(block.content.headerRows);
 
   return (
     <Components.Generic.Menu.Item
       className={"bn-menu-item"}
       checked={isHeaderRow}
       onClick={() => {
-        // The block may have been modified and out of date, so we get the latest block
-        const block = editor.getBlock(props.block.id);
-        if (!block) {
-          return;
-        }
         editor.updateBlock(block, {
           ...block,
           content: {
             ...block.content,
             headerRows: isHeaderRow ? undefined : 1,
-          },
+          } as any,
         });
       }}
     >
@@ -66,9 +73,9 @@ export const TableHeaderRowButton = <
 export const TableHeaderColumnButton = <
   I extends InlineContentSchema = DefaultInlineContentSchema,
   S extends StyleSchema = DefaultStyleSchema,
->(
-  props: TableHandleMenuProps<I, S> & { orientation: "row" | "column" },
-) => {
+>(props: {
+  orientation: "row" | "column";
+}) => {
   const Components = useComponentsContext()!;
   const dict = useDictionary();
 
@@ -77,11 +84,19 @@ export const TableHeaderColumnButton = <
     I,
     S
   >();
-  const tableHandles = editor.tableHandles;
+  const tableHandles = useExtension(TableHandles);
+  const block = useExtensionState(TableHandles, {
+    selector: (state) => state?.block,
+  });
+  const index = useExtensionState(TableHandles, {
+    selector: (state) =>
+      props.orientation === "column" ? state?.colIndex : state?.rowIndex,
+  });
 
   if (
     !tableHandles ||
-    props.index !== 0 ||
+    index !== 0 ||
+    !block ||
     props.orientation !== "column" ||
     !editor.settings.tables.headers
   ) {
@@ -89,24 +104,19 @@ export const TableHeaderColumnButton = <
   }
 
   // We only support 1 header column for now
-  const isHeaderColumn = Boolean(props.block.content.headerCols);
+  const isHeaderColumn = Boolean(block.content.headerCols);
 
   return (
     <Components.Generic.Menu.Item
       className={"bn-menu-item"}
       checked={isHeaderColumn}
       onClick={() => {
-        // The block may have been modified and out of date, so we get the latest block
-        const block = editor.getBlock(props.block.id);
-        if (!block) {
-          return;
-        }
         editor.updateBlock(block, {
           ...block,
           content: {
             ...block.content,
             headerCols: isHeaderColumn ? undefined : 1,
-          },
+          } as any,
         });
       }}
     >
diff --git a/packages/react/src/components/TableHandles/TableHandleMenu/TableHandleMenu.tsx b/packages/react/src/components/TableHandles/TableHandleMenu/TableHandleMenu.tsx
index 3366660bc4..ad90342261 100644
--- a/packages/react/src/components/TableHandles/TableHandleMenu/TableHandleMenu.tsx
+++ b/packages/react/src/components/TableHandles/TableHandleMenu/TableHandleMenu.tsx
@@ -1,24 +1,15 @@
-import {
-  DefaultInlineContentSchema,
-  DefaultStyleSchema,
-  InlineContentSchema,
-  StyleSchema,
-} from "@blocknote/core";
 import { ReactNode } from "react";
 
 import { useComponentsContext } from "../../../editor/ComponentsContext.js";
 import { AddButton } from "./DefaultButtons/AddButton.js";
 import { DeleteButton } from "./DefaultButtons/DeleteButton.js";
-import { TableHandleMenuProps } from "./TableHandleMenuProps.js";
 import { ColorPickerButton } from "./DefaultButtons/ColorPicker.js";
 import { TableHeaderColumnButton } from "./DefaultButtons/TableHeaderButton.js";
 import { TableHeaderRowButton } from "./DefaultButtons/TableHeaderButton.js";
+import { TableHandleMenuProps } from "./TableHandleMenuProps.js";
 
-export const TableHandleMenu = <
-  I extends InlineContentSchema = DefaultInlineContentSchema,
-  S extends StyleSchema = DefaultStyleSchema,
->(
-  props: TableHandleMenuProps<I, S> & { children?: ReactNode },
+export const TableHandleMenu = (
+  props: TableHandleMenuProps & { children?: ReactNode },
 ) => {
   const Components = useComponentsContext()!;
 
@@ -26,38 +17,18 @@ export const TableHandleMenu = <
     <Components.Generic.Menu.Dropdown className={"bn-table-handle-menu"}>
       {props.children || (
         <>
-          <DeleteButton
-            orientation={props.orientation}
-            block={props.block}
-            index={props.index}
-          />
+          <DeleteButton orientation={props.orientation} />
           <AddButton
             orientation={props.orientation}
-            block={props.block}
-            index={props.index}
             side={props.orientation === "row" ? "above" : ("left" as any)}
           />
           <AddButton
             orientation={props.orientation}
-            block={props.block}
-            index={props.index}
             side={props.orientation === "row" ? "below" : ("right" as any)}
           />
-          <TableHeaderRowButton
-            orientation={props.orientation}
-            block={props.block}
-            index={props.index}
-          />
-          <TableHeaderColumnButton
-            orientation={props.orientation}
-            block={props.block}
-            index={props.index}
-          />
-          <ColorPickerButton
-            orientation={props.orientation}
-            block={props.block}
-            index={props.index}
-          />
+          <TableHeaderRowButton orientation={props.orientation} />
+          <TableHeaderColumnButton orientation={props.orientation} />
+          <ColorPickerButton orientation={props.orientation} />
         </>
       )}
     </Components.Generic.Menu.Dropdown>
diff --git a/packages/react/src/components/TableHandles/TableHandleMenu/TableHandleMenuProps.ts b/packages/react/src/components/TableHandles/TableHandleMenu/TableHandleMenuProps.ts
index 74477f8286..b064bd1869 100644
--- a/packages/react/src/components/TableHandles/TableHandleMenu/TableHandleMenuProps.ts
+++ b/packages/react/src/components/TableHandles/TableHandleMenu/TableHandleMenuProps.ts
@@ -1,17 +1,3 @@
-import {
-  DefaultBlockSchema,
-  DefaultInlineContentSchema,
-  DefaultStyleSchema,
-  InlineContentSchema,
-  SpecificBlock,
-  StyleSchema,
-} from "@blocknote/core";
-
-export type TableHandleMenuProps<
-  I extends InlineContentSchema = DefaultInlineContentSchema,
-  S extends StyleSchema = DefaultStyleSchema,
-> = {
+export type TableHandleMenuProps = {
   orientation: "row" | "column";
-  block: SpecificBlock<{ table: DefaultBlockSchema["table"] }, "table", I, S>;
-  index: number;
 };
diff --git a/packages/react/src/components/TableHandles/TableHandleProps.ts b/packages/react/src/components/TableHandles/TableHandleProps.ts
index 924e77be94..4a2fc65ac6 100644
--- a/packages/react/src/components/TableHandles/TableHandleProps.ts
+++ b/packages/react/src/components/TableHandles/TableHandleProps.ts
@@ -1,53 +1,7 @@
-import {
-  BlockNoteEditor,
-  DefaultBlockSchema,
-  DefaultInlineContentSchema,
-  DefaultStyleSchema,
-  InlineContentSchema,
-  StyleSchema,
-  TableHandlesState,
-} from "@blocknote/core";
-import { DragEvent, FC } from "react";
+import { FC } from "react";
 
-import { DragHandleMenuProps } from "../SideMenu/DragHandleMenu/DragHandleMenuProps.js";
-
-export type TableHandleProps<
-  I extends InlineContentSchema = DefaultInlineContentSchema,
-  S extends StyleSchema = DefaultStyleSchema,
-> = {
-  editor: BlockNoteEditor<
-    {
-      table: DefaultBlockSchema["table"];
-    },
-    I,
-    S
-  >;
+export type TableHandleProps = {
   orientation: "row" | "column";
-  index: number;
-  dragStart: (e: DragEvent) => void;
-  showOtherSide: () => void;
-  hideOtherSide: () => void;
-  menuContainer: HTMLDivElement;
-  tableHandleMenu?: FC<
-    DragHandleMenuProps<
-      {
-        table: DefaultBlockSchema["table"];
-      },
-      I,
-      S
-    >
-  >;
-} & Pick<TableHandlesState<I, S>, "block"> &
-  Pick<
-    Exclude<
-      BlockNoteEditor<
-        {
-          table: DefaultBlockSchema["table"];
-        },
-        I,
-        S
-      >["tableHandles"],
-      undefined
-    >,
-    "dragEnd" | "freezeHandles" | "unfreezeHandles"
-  >;
+  hideOtherElements: (hide: boolean) => void;
+  tableHandleMenu?: FC;
+};
diff --git a/packages/react/src/components/TableHandles/TableHandlesController.tsx b/packages/react/src/components/TableHandles/TableHandlesController.tsx
index 51f6d7626e..6a39179fb1 100644
--- a/packages/react/src/components/TableHandles/TableHandlesController.tsx
+++ b/packages/react/src/components/TableHandles/TableHandlesController.tsx
@@ -2,98 +2,121 @@ import {
   BlockSchema,
   DefaultInlineContentSchema,
   DefaultStyleSchema,
+  getNodeById,
   InlineContentSchema,
   StyleSchema,
 } from "@blocknote/core";
-import { FC, useCallback, useMemo, useState } from "react";
+import { TableHandles } from "@blocknote/core/extensions";
+import { FC, useMemo, useState } from "react";
 
-import { FloatingPortal } from "@floating-ui/react";
+import { offset, size, VirtualElement } from "@floating-ui/react";
 import { useBlockNoteEditor } from "../../hooks/useBlockNoteEditor.js";
-import { useUIPluginState } from "../../hooks/useUIPluginState.js";
 import { ExtendButton } from "./ExtendButton/ExtendButton.js";
 import { ExtendButtonProps } from "./ExtendButton/ExtendButtonProps.js";
 import { TableHandle } from "./TableHandle.js";
-import { TableHandleProps } from "./TableHandleProps.js";
-import { useExtendButtonsPositioning } from "./hooks/useExtendButtonsPositioning.js";
-import { useTableHandlesPositioning } from "./hooks/useTableHandlesPositioning.js";
 import { TableCellButton } from "./TableCellButton.js";
 import { TableCellButtonProps } from "./TableCellButtonProps.js";
+import { useExtensionState } from "../../hooks/useExtension.js";
+import { TableHandleProps } from "./TableHandleProps.js";
+import { GenericPopover } from "../Popovers/GenericPopover.js";
 
 export const TableHandlesController = <
   I extends InlineContentSchema = DefaultInlineContentSchema,
   S extends StyleSchema = DefaultStyleSchema,
 >(props: {
-  tableCellHandle?: FC<TableCellButtonProps<I, S>>;
-  tableHandle?: FC<TableHandleProps<I, S>>;
-  extendButton?: FC<ExtendButtonProps<I, S>>;
+  tableCellHandle?: FC<TableCellButtonProps>;
+  tableHandle?: FC<TableHandleProps>;
+  extendButton?: FC<ExtendButtonProps>;
 }) => {
   const editor = useBlockNoteEditor<BlockSchema, I, S>();
 
-  const [menuContainerRef, setMenuContainerRef] =
-    useState<HTMLDivElement | null>(null);
+  const [onlyShownElement, setOnlyShownElement] = useState<
+    | "rowTableHandle"
+    | "columnTableHandle"
+    | "tableCellHandle"
+    | "extendRowsButton"
+    | "extendColumnsButton"
+    | undefined
+  >();
+
+  const state = useExtensionState(TableHandles);
+
+  const elements = useMemo(() => {
+    if (
+      state === undefined ||
+      state.rowIndex === undefined ||
+      state.colIndex === undefined
+    ) {
+      return {};
+    }
 
-  if (!editor.tableHandles) {
-    throw new Error(
-      "TableHandlesController can only be used when BlockNote editor schema contains table block",
+    // TODO use the location API for this
+    const nodePosInfo = getNodeById(
+      state.block.id,
+      editor.prosemirrorState.doc,
     );
-  }
+    if (!nodePosInfo) {
+      return {};
+    }
 
-  const callbacks = {
-    rowDragStart: editor.tableHandles.rowDragStart,
-    colDragStart: editor.tableHandles.colDragStart,
-    dragEnd: editor.tableHandles.dragEnd,
-    freezeHandles: editor.tableHandles.freezeHandles,
-    unfreezeHandles: editor.tableHandles.unfreezeHandles,
-  };
-
-  const { freezeHandles, unfreezeHandles } = callbacks;
-
-  const onStartExtend = useCallback(() => {
-    freezeHandles();
-    setHideCol(true);
-    setHideRow(true);
-  }, [freezeHandles]);
-
-  const onEndExtend = useCallback(() => {
-    unfreezeHandles();
-    setHideCol(false);
-    setHideRow(false);
-  }, [unfreezeHandles]);
-
-  const state = useUIPluginState(
-    editor.tableHandles.onUpdate.bind(editor.tableHandles),
-  );
+    const tableBeforePos = nodePosInfo.posBeforeNode + 1;
 
-  const draggingState = useMemo(() => {
-    return state?.draggingState
-      ? {
-          draggedCellOrientation: state?.draggingState?.draggedCellOrientation,
-          mousePos: state?.draggingState?.mousePos,
-        }
-      : undefined;
-    // eslint-disable-next-line react-hooks/exhaustive-deps
-  }, [
-    state?.draggingState,
-    state?.draggingState?.draggedCellOrientation,
-    state?.draggingState?.mousePos,
-  ]);
-
-  const { rowHandle, colHandle, cellHandle } = useTableHandlesPositioning(
-    state?.show || false,
-    state?.referencePosCell || null,
-    state?.referencePosTable || null,
-    draggingState,
-  );
+    const tableElement = editor.prosemirrorView.domAtPos(
+      tableBeforePos + 1,
+    ).node;
+    if (!(tableElement instanceof Element)) {
+      return undefined;
+    }
 
-  const { addOrRemoveColumnsButton, addOrRemoveRowsButton } =
-    useExtendButtonsPositioning(
-      state?.showAddOrRemoveColumnsButton || false,
-      state?.showAddOrRemoveRowsButton || false,
-      state?.referencePosTable || null,
-    );
+    const rowBeforePos = editor.prosemirrorState.doc
+      .resolve(tableBeforePos + 1)
+      .posAtIndex(state.rowIndex);
+    const cellBeforePos = editor.prosemirrorState.doc
+      .resolve(rowBeforePos + 1)
+      .posAtIndex(state.colIndex);
 
-  const [hideRow, setHideRow] = useState<boolean>(false);
-  const [hideCol, setHideCol] = useState<boolean>(false);
+    const cellElement = editor.prosemirrorView.domAtPos(cellBeforePos + 1).node;
+    if (!(cellElement instanceof Element)) {
+      return undefined;
+    }
+
+    const tableBoundingRect = tableElement.getBoundingClientRect();
+    const cellBoundingRect = cellElement.getBoundingClientRect();
+
+    const rowVirtualElement: VirtualElement = {
+      getBoundingClientRect: () =>
+        new DOMRect(
+          tableBoundingRect.x,
+          state.draggingState &&
+          state.draggingState.draggedCellOrientation === "row"
+            ? state.draggingState.mousePos - cellBoundingRect.height / 2
+            : cellBoundingRect.y,
+          tableBoundingRect.width,
+          cellBoundingRect.height,
+        ),
+      contextElement: tableElement,
+    };
+    const columnVirtualElement: VirtualElement = {
+      getBoundingClientRect: () =>
+        new DOMRect(
+          state.draggingState &&
+          state.draggingState.draggedCellOrientation === "col"
+            ? state.draggingState.mousePos - cellBoundingRect.width / 2
+            : cellBoundingRect.x,
+          tableBoundingRect.y,
+          cellBoundingRect.width,
+          tableBoundingRect.height,
+        ),
+      contextElement: tableElement,
+    };
+
+    return {
+      tableElement,
+      cellElement,
+      rowVirtualElement,
+      columnVirtualElement,
+    };
+  }, [editor, state]);
 
   if (!state) {
     return null;
@@ -105,97 +128,158 @@ export const TableHandlesController = <
 
   return (
     <>
-      <div ref={setMenuContainerRef}></div>
-      {/* we want to make sure the elements are clipped by the .tableWrapper element (so that we scroll the table, widgets also dissappear)
-      we do this by rendering in a portal into the table's widget container (defined in TableBlockContent.ts)
-      */}
-      <FloatingPortal root={state.widgetContainer}>
-        {!hideRow &&
-          menuContainerRef &&
-          rowHandle.isMounted &&
-          state.rowIndex !== undefined && (
-            <div ref={rowHandle.ref} style={rowHandle.style}>
-              <TableHandleComponent
-                // This "as any" unfortunately seems complicated to fix
-                editor={editor as any}
-                orientation={"row"}
-                showOtherSide={() => setHideCol(false)}
-                hideOtherSide={() => setHideCol(true)}
-                index={state.rowIndex}
-                block={state.block}
-                dragStart={callbacks.rowDragStart}
-                dragEnd={callbacks.dragEnd}
-                freezeHandles={callbacks.freezeHandles}
-                unfreezeHandles={callbacks.unfreezeHandles}
-                menuContainer={menuContainerRef}
-              />
-            </div>
+      <GenericPopover
+        reference={elements?.rowVirtualElement}
+        useFloatingOptions={{
+          open:
+            state.show &&
+            state.rowIndex !== undefined &&
+            (!onlyShownElement || onlyShownElement === "rowTableHandle"),
+          placement: "left",
+          middleware: [offset(-10)],
+        }}
+        elementProps={{
+          style: {
+            zIndex: 10,
+          },
+        }}
+      >
+        {state.show &&
+          state.rowIndex !== undefined &&
+          (!onlyShownElement || onlyShownElement === "rowTableHandle") && (
+            <TableHandleComponent
+              orientation="row"
+              hideOtherElements={(hide) =>
+                setOnlyShownElement(hide ? "rowTableHandle" : undefined)
+              }
+            />
           )}
-        {!hideCol &&
-          menuContainerRef &&
-          colHandle.isMounted &&
-          state.colIndex !== undefined && (
-            <div ref={colHandle.ref} style={colHandle.style}>
-              <TableHandleComponent
-                editor={editor as any}
-                orientation={"column"}
-                showOtherSide={() => setHideRow(false)}
-                hideOtherSide={() => setHideRow(true)}
-                index={state.colIndex}
-                block={state.block}
-                dragStart={callbacks.colDragStart}
-                dragEnd={callbacks.dragEnd}
-                freezeHandles={callbacks.freezeHandles}
-                unfreezeHandles={callbacks.unfreezeHandles}
-                menuContainer={menuContainerRef}
-              />
-            </div>
+      </GenericPopover>
+      <GenericPopover
+        reference={elements?.columnVirtualElement}
+        useFloatingOptions={{
+          open:
+            state.show &&
+            state.colIndex !== undefined &&
+            (!onlyShownElement || onlyShownElement === "columnTableHandle"),
+          placement: "top",
+          middleware: [offset(-12)],
+        }}
+        elementProps={{
+          style: {
+            zIndex: 10,
+          },
+        }}
+      >
+        {state.show &&
+          state.colIndex !== undefined &&
+          (!onlyShownElement || onlyShownElement === "columnTableHandle") && (
+            <TableHandleComponent
+              orientation="column"
+              hideOtherElements={(hide) =>
+                setOnlyShownElement(hide ? "columnTableHandle" : undefined)
+              }
+            />
           )}
-
-        {menuContainerRef &&
-          cellHandle.isMounted &&
+      </GenericPopover>
+      <GenericPopover
+        reference={elements?.cellElement}
+        useFloatingOptions={{
+          open:
+            state.show &&
+            state.rowIndex !== undefined &&
+            state.colIndex !== undefined &&
+            (!onlyShownElement || onlyShownElement === "tableCellHandle"),
+          placement: "top-end",
+          middleware: [offset({ mainAxis: -15, crossAxis: -1 })],
+        }}
+        elementProps={{
+          style: {
+            zIndex: 10,
+          },
+        }}
+      >
+        {state.show &&
+          state.rowIndex !== undefined &&
           state.colIndex !== undefined &&
-          state.rowIndex !== undefined && (
-            <div ref={cellHandle.ref} style={cellHandle.style}>
-              <TableCellHandleComponent
-                editor={editor as any}
-                block={state.block}
-                rowIndex={state.rowIndex}
-                colIndex={state.colIndex}
-                menuContainer={menuContainerRef}
-                freezeHandles={callbacks.freezeHandles}
-                unfreezeHandles={callbacks.unfreezeHandles}
-              />
-            </div>
+          (!onlyShownElement || onlyShownElement === "tableCellHandle") && (
+            <TableCellHandleComponent
+              hideOtherElements={(hide) =>
+                setOnlyShownElement(hide ? "tableCellHandle" : undefined)
+              }
+            />
           )}
-
-        {/* note that the extend buttons are always shown (we don't look at isMounted etc, 
-        because otherwise the table slightly shifts when they unmount  */}
-        <div
-          ref={addOrRemoveRowsButton.ref}
-          style={addOrRemoveRowsButton.style}
-        >
-          <ExtendButtonComponent
-            editor={editor as any}
-            orientation={"addOrRemoveRows"}
-            block={state.block}
-            onMouseDown={onStartExtend}
-            onMouseUp={onEndExtend}
-          />
-        </div>
-        <div
-          ref={addOrRemoveColumnsButton.ref}
-          style={addOrRemoveColumnsButton.style}
-        >
-          <ExtendButtonComponent
-            editor={editor as any}
-            orientation={"addOrRemoveColumns"}
-            block={state.block}
-            onMouseDown={onStartExtend}
-            onMouseUp={onEndExtend}
-          />
-        </div>
-      </FloatingPortal>
+      </GenericPopover>
+      <GenericPopover
+        reference={elements?.tableElement}
+        useFloatingOptions={{
+          open:
+            state.show &&
+            state.showAddOrRemoveRowsButton &&
+            (!onlyShownElement || onlyShownElement === "extendRowsButton"),
+          placement: "bottom",
+          middleware: [
+            size({
+              apply({ rects, elements }) {
+                Object.assign(elements.floating.style, {
+                  width: `${rects.reference.width}px`,
+                });
+              },
+            }),
+          ],
+        }}
+        elementProps={{
+          style: {
+            zIndex: 10,
+          },
+        }}
+      >
+        {state.show &&
+          state.showAddOrRemoveRowsButton &&
+          (!onlyShownElement || onlyShownElement === "extendRowsButton") && (
+            <ExtendButtonComponent
+              orientation="addOrRemoveRows"
+              hideOtherElements={(hide) =>
+                setOnlyShownElement(hide ? "extendRowsButton" : undefined)
+              }
+            />
+          )}
+      </GenericPopover>
+      <GenericPopover
+        reference={elements?.tableElement}
+        useFloatingOptions={{
+          open:
+            state.show &&
+            state.showAddOrRemoveColumnsButton &&
+            (!onlyShownElement || onlyShownElement === "extendColumnsButton"),
+          placement: "right",
+          middleware: [
+            size({
+              apply({ rects, elements }) {
+                Object.assign(elements.floating.style, {
+                  height: `${rects.reference.height}px`,
+                });
+              },
+            }),
+          ],
+        }}
+        elementProps={{
+          style: {
+            zIndex: 10,
+          },
+        }}
+      >
+        {state.show &&
+          state.showAddOrRemoveColumnsButton &&
+          (!onlyShownElement || onlyShownElement === "extendColumnsButton") && (
+            <ExtendButtonComponent
+              orientation="addOrRemoveColumns"
+              hideOtherElements={(hide) =>
+                setOnlyShownElement(hide ? "extendColumnsButton" : undefined)
+              }
+            />
+          )}
+      </GenericPopover>
     </>
   );
 };
diff --git a/packages/react/src/editor/BlockNoteDefaultUI.tsx b/packages/react/src/editor/BlockNoteDefaultUI.tsx
index 2f6fd28cf9..55f2f71332 100644
--- a/packages/react/src/editor/BlockNoteDefaultUI.tsx
+++ b/packages/react/src/editor/BlockNoteDefaultUI.tsx
@@ -1,3 +1,11 @@
+import {
+  Comments,
+  FormattingToolbar,
+  SideMenu,
+  SuggestionMenu,
+  TableHandles,
+} from "@blocknote/core/extensions";
+
 import { FloatingComposerController } from "../components/Comments/FloatingComposerController.js";
 import { FloatingThreadController } from "../components/Comments/FloatingThreadController.js";
 import { FilePanelController } from "../components/FilePanel/FilePanelController.js";
@@ -70,24 +78,27 @@ export function BlockNoteDefaultUI(props: BlockNoteDefaultUIProps) {
 
   return (
     <>
-      {props.formattingToolbar !== false && <FormattingToolbarController />}
+      {editor.getExtension(FormattingToolbar) &&
+        props.formattingToolbar !== false && <FormattingToolbarController />}
       {props.linkToolbar !== false && <LinkToolbarController />}
-      {props.slashMenu !== false && (
+      {editor.getExtension(SuggestionMenu) && props.slashMenu !== false && (
         <SuggestionMenuController triggerCharacter="/" />
       )}
-      {props.emojiPicker !== false && (
+      {editor.getExtension(SuggestionMenu) && props.emojiPicker !== false && (
         <GridSuggestionMenuController
           triggerCharacter=":"
           columns={10}
           minQueryLength={2}
         />
       )}
-      {props.sideMenu !== false && <SideMenuController />}
-      {editor.filePanel && props.filePanel !== false && <FilePanelController />}
-      {editor.tableHandles && props.tableHandles !== false && (
+      {editor.getExtension(SideMenu) && props.sideMenu !== false && (
+        <SideMenuController />
+      )}
+      {props.filePanel !== false && <FilePanelController />}
+      {editor.getExtension(TableHandles) && props.tableHandles !== false && (
         <TableHandlesController />
       )}
-      {editor.comments && props.comments !== false && (
+      {editor.getExtension(Comments) && props.comments !== false && (
         <>
           <FloatingComposerController />
           <FloatingThreadController />
diff --git a/packages/react/src/editor/styles.css b/packages/react/src/editor/styles.css
index 1cd909599f..3862946986 100644
--- a/packages/react/src/editor/styles.css
+++ b/packages/react/src/editor/styles.css
@@ -230,7 +230,14 @@
   background-color: var(--bn-colors-highlights-pink-background);
 }
 
-/* Matching Side Menu height to block line height */
+/* Sets base Z-index on UI elements, i.e. elements which use the 
+`GenericPopover` component. If a z-index is passed to the `GenericPopover` via 
+inline styles, it is added to the base z-index. */
+.bn-container {
+  --bn-ui-base-z-index: 0;
+}
+
+/* Matches Side Menu height to block line height */
 .bn-side-menu {
   height: 30px;
 }
diff --git a/packages/react/src/hooks/useActiveStyles.ts b/packages/react/src/hooks/useActiveStyles.ts
index 0b0e97333d..a539d50f71 100644
--- a/packages/react/src/hooks/useActiveStyles.ts
+++ b/packages/react/src/hooks/useActiveStyles.ts
@@ -1,13 +1,11 @@
-import { BlockNoteEditor, StyleSchema } from "@blocknote/core";
+import { BlockNoteEditor, Styles, StyleSchema } from "@blocknote/core";
 
-import { useCallback, useState } from "react";
 import { useBlockNoteContext } from "../editor/BlockNoteContext.js";
-import { useEditorChange } from "./useEditorChange.js";
-import { useEditorSelectionChange } from "./useEditorSelectionChange.js";
+import { useEditorState } from "./useEditorState.js";
 
 export function useActiveStyles<T extends StyleSchema>(
   editor?: BlockNoteEditor<any, any, T>,
-) {
+): Styles<T> {
   const editorContext = useBlockNoteContext<any, any, T>();
   if (!editor) {
     editor = editorContext?.editor;
@@ -19,19 +17,10 @@ export function useActiveStyles<T extends StyleSchema>(
     );
   }
 
-  const e = editor;
-
-  const [styles, setStyles] = useState(() => e.getActiveStyles());
-
-  const updateStyles = useCallback(() => {
-    setStyles(e.getActiveStyles());
-  }, [e]);
-
-  // Updates state on editor content change.
-  useEditorChange(updateStyles, e);
-
-  // Updates state on selection change.
-  useEditorSelectionChange(updateStyles, e);
+  const styles = useEditorState({
+    editor,
+    selector: ({ editor }) => editor.getActiveStyles(),
+  });
 
   return styles;
 }
diff --git a/packages/react/src/hooks/useEditorSelectionBoundingBox.ts b/packages/react/src/hooks/useEditorSelectionBoundingBox.ts
index 8b7c5184e0..12200e0b2a 100644
--- a/packages/react/src/hooks/useEditorSelectionBoundingBox.ts
+++ b/packages/react/src/hooks/useEditorSelectionBoundingBox.ts
@@ -1,43 +1,13 @@
 import type { BlockNoteEditor } from "@blocknote/core";
-import { useCallback, useEffect, useState } from "react";
-import { useBlockNoteEditor } from "./useBlockNoteEditor.js";
-import { useEditorSelectionChange } from "./useEditorSelectionChange.js";
+import { useEditorState } from "./useEditorState.js";
 
 export function useEditorSelectionBoundingBox(
   enabled?: boolean,
   editor?: BlockNoteEditor<any, any, any>,
 ) {
-  const ctxEditor = useBlockNoteEditor();
-  editor = editor || ctxEditor;
-
-  const [boundingBox, setBoundingBox] = useState(() => {
-    if (!enabled) {
-      return undefined;
-    }
-
-    return editor.getSelectionBoundingBox();
+  return useEditorState({
+    editor,
+    selector: ({ editor }) =>
+      enabled ? editor.getSelectionBoundingBox() : undefined,
   });
-
-  const cb = useCallback(() => {
-    if (!enabled) {
-      return;
-    }
-
-    const selection = editor.getSelectionBoundingBox();
-
-    setBoundingBox(selection);
-  }, [editor, enabled]);
-
-  useEditorSelectionChange(cb, editor, true);
-
-  // initial value when `enabled` changes
-  useEffect(() => {
-    if (!enabled) {
-      setBoundingBox(undefined);
-    } else {
-      setBoundingBox(editor.getSelectionBoundingBox());
-    }
-  }, [enabled, editor]);
-
-  return boundingBox;
 }
diff --git a/packages/react/src/hooks/useEditorState.ts b/packages/react/src/hooks/useEditorState.ts
new file mode 100644
index 0000000000..f83897bc1b
--- /dev/null
+++ b/packages/react/src/hooks/useEditorState.ts
@@ -0,0 +1,229 @@
+import type { BlockNoteEditor } from "@blocknote/core";
+import deepEqual from "fast-deep-equal/es6/react";
+import { useDebugValue, useEffect, useLayoutEffect, useState } from "react";
+import { useSyncExternalStoreWithSelector } from "use-sync-external-store/shim/with-selector";
+import { useBlockNoteContext } from "../editor/BlockNoteContext.js";
+
+const useIsomorphicLayoutEffect =
+  typeof window !== "undefined" ? useLayoutEffect : useEffect;
+
+export type EditorStateSnapshot<
+  TEditor extends BlockNoteEditor<any, any, any> | null = BlockNoteEditor<
+    any,
+    any,
+    any
+  > | null,
+> = {
+  editor: TEditor;
+  transactionNumber: number;
+};
+
+export type UseEditorStateOptions<
+  TSelectorResult,
+  TEditor extends BlockNoteEditor<any, any, any> | null = BlockNoteEditor<
+    any,
+    any,
+    any
+  > | null,
+> = {
+  /**
+   * The editor instance. If not provided, will use the editor from BlockNoteContext.
+   */
+  editor?: TEditor;
+
+  /**
+   * A selector function to determine the value to compare for re-rendering.
+   */
+  selector: (context: EditorStateSnapshot<TEditor>) => TSelectorResult;
+
+  /**
+   * A custom equality function to determine if the editor should re-render.
+   * @default `deepEqual` from `fast-deep-equal`
+   */
+  equalityFn?: (a: TSelectorResult, b: TSelectorResult | null) => boolean;
+
+  /**
+   * The event to subscribe to.
+   * @default "all"
+   */
+  on?: "all" | "selection" | "change";
+};
+
+/**
+ * To synchronize the editor instance with the component state,
+ * we need to create a separate instance that is not affected by the component re-renders.
+ */
+class EditorStateManager<
+  TEditor extends BlockNoteEditor<any, any, any> | null = BlockNoteEditor<
+    any,
+    any,
+    any
+  > | null,
+> {
+  private transactionNumber = 0;
+
+  private lastTransactionNumber = 0;
+
+  private lastSnapshot: EditorStateSnapshot<TEditor>;
+
+  private editor: TEditor;
+
+  private subscribers = new Set<() => void>();
+
+  constructor(initialEditor: TEditor) {
+    this.editor = initialEditor;
+    this.lastSnapshot = { editor: initialEditor, transactionNumber: 0 };
+
+    this.getSnapshot = this.getSnapshot.bind(this);
+    this.getServerSnapshot = this.getServerSnapshot.bind(this);
+    this.watch = this.watch.bind(this);
+    this.subscribe = this.subscribe.bind(this);
+  }
+
+  /**
+   * Get the current editor instance.
+   */
+  getSnapshot(): EditorStateSnapshot<TEditor> {
+    if (this.transactionNumber === this.lastTransactionNumber) {
+      return this.lastSnapshot;
+    }
+    this.lastTransactionNumber = this.transactionNumber;
+    this.lastSnapshot = {
+      editor: this.editor,
+      transactionNumber: this.transactionNumber,
+    };
+    return this.lastSnapshot;
+  }
+
+  /**
+   * Always disable the editor on the server-side.
+   */
+  getServerSnapshot(): EditorStateSnapshot<null> {
+    return { editor: null, transactionNumber: 0 };
+  }
+
+  /**
+   * Subscribe to the editor instance's changes.
+   */
+  subscribe(callback: () => void): () => void {
+    this.subscribers.add(callback);
+    return () => {
+      this.subscribers.delete(callback);
+    };
+  }
+
+  /**
+   * Watch the editor instance for changes.
+   */
+  watch(
+    nextEditor: BlockNoteEditor<any, any, any> | null,
+    on: "all" | "selection" | "change",
+  ): undefined | (() => void) {
+    this.editor = nextEditor as TEditor;
+
+    if (this.editor) {
+      /**
+       * This will force a re-render when the editor state changes.
+       * This is to support things like `editor.can().toggleBold()` in components that `useEditor`.
+       * This could be more efficient, but it's a good trade-off for now.
+       */
+      const fn = () => {
+        this.transactionNumber += 1;
+        this.subscribers.forEach((callback) => callback());
+      };
+
+      const currentTiptapEditor = this.editor._tiptapEditor;
+
+      const EVENT_TYPES = {
+        all: "transaction",
+        selection: "selectionUpdate",
+        change: "update",
+      } as const;
+
+      currentTiptapEditor.on(EVENT_TYPES[on], fn);
+      return () => {
+        currentTiptapEditor.off(EVENT_TYPES[on], fn);
+      };
+    }
+
+    return undefined;
+  }
+}
+
+/**
+ * This hook allows you to watch for changes on the editor instance.
+ * It will allow you to select a part of the editor state and re-render the component when it changes.
+ * @example
+ * ```tsx
+ * const { currentSelection } = useEditorState({
+ *  selector: snapshot => ({ currentSelection: snapshot.editor?._tiptapEditor.state.selection }),
+ * })
+ * ```
+ */
+export function useEditorState<TSelectorResult>(
+  options: UseEditorStateOptions<
+    TSelectorResult,
+    BlockNoteEditor<any, any, any>
+  >,
+): TSelectorResult;
+/**
+ * This hook allows you to watch for changes on the editor instance.
+ * It will allow you to select a part of the editor state and re-render the component when it changes.
+ * @example
+ * ```tsx
+ * const { currentSelection } = useEditorState({
+ *  selector: snapshot => ({ currentSelection: snapshot.editor?._tiptapEditor.state.selection }),
+ * })
+ * ```
+ */
+export function useEditorState<TSelectorResult>(
+  options: UseEditorStateOptions<
+    TSelectorResult,
+    BlockNoteEditor<any, any, any> | null
+  >,
+): TSelectorResult | null;
+
+/**
+ * This hook allows you to watch for changes on the editor instance.
+ * It will allow you to select a part of the editor state and re-render the component when it changes.
+ * @example
+ * ```tsx
+ * const { currentSelection } = useEditorState({
+ *  selector: snapshot => ({ currentSelection: snapshot.editor?._tiptapEditor.state.selection }),
+ * })
+ * ```
+ */
+export function useEditorState<TSelectorResult>(
+  options:
+    | UseEditorStateOptions<TSelectorResult, BlockNoteEditor<any, any, any>>
+    | UseEditorStateOptions<
+        TSelectorResult,
+        BlockNoteEditor<any, any, any> | null
+      >,
+): TSelectorResult | null {
+  const editorContext = useBlockNoteContext();
+  const editor = options.editor || editorContext?.editor || null;
+  const on = options.on || "all";
+
+  const [editorStateManager] = useState(() => new EditorStateManager(editor));
+
+  // Using the `useSyncExternalStore` hook to sync the editor instance with the component state
+  const selectedState = useSyncExternalStoreWithSelector(
+    editorStateManager.subscribe,
+    editorStateManager.getSnapshot,
+    editorStateManager.getServerSnapshot,
+    options.selector as UseEditorStateOptions<
+      TSelectorResult,
+      BlockNoteEditor<any, any, any> | null
+    >["selector"],
+    options.equalityFn ?? deepEqual,
+  );
+
+  useIsomorphicLayoutEffect(() => {
+    return editorStateManager.watch(editor, on);
+  }, [editor, editorStateManager, on]);
+
+  useDebugValue(selectedState);
+
+  return selectedState;
+}
diff --git a/packages/react/src/hooks/useElementHover.ts b/packages/react/src/hooks/useElementHover.ts
new file mode 100644
index 0000000000..c84acc9c7f
--- /dev/null
+++ b/packages/react/src/hooks/useElementHover.ts
@@ -0,0 +1,433 @@
+import * as React from "react";
+import { isElement } from "@floating-ui/utils/dom";
+import {
+  isMouseLikePointerType,
+  useLatestRef,
+  useEffectEvent,
+  useModernLayoutEffect,
+} from "@floating-ui/react/utils";
+import type {
+  Delay,
+  ElementProps,
+  FloatingContext,
+  OpenChangeReason,
+} from "@floating-ui/react";
+
+function clearTimeoutIfSet(timeoutRef: React.RefObject<number>) {
+  if (timeoutRef.current !== -1) {
+    clearTimeout(timeoutRef.current);
+    timeoutRef.current = -1;
+  }
+}
+
+export interface UseElementHoverProps<T = HTMLElement> {
+  /**
+   * Whether the Hook is enabled, including all internal Effects and event
+   * handlers.
+   * @default true
+   */
+  enabled?: boolean;
+  /**
+   * Waits for the specified time when the event listener runs before changing
+   * the `open` state.
+   * @default 0
+   */
+  delay?: Delay | (() => Delay);
+  /**
+   * Waits until the user's cursor is at "rest" over the reference element
+   * before changing the `open` state.
+   * @default 0
+   */
+  restMs?: number | (() => number);
+  /**
+   * Whether the logic only runs for mouse input, ignoring touch input.
+   * Note: due to a bug with Linux Chrome, "pen" inputs are considered "mouse".
+   * @default false
+   */
+  mouseOnly?: boolean;
+  /**
+   * The element to attach mouse event listeners to.
+   */
+  attachTo: () => HTMLElement;
+  /**
+   * Function to get the element at the hover position. Returns null to disable hover.
+   */
+  getElementAtHover: (target: EventTarget) => T | null;
+  /**
+   * Callback when an element is hovered.
+   */
+  onHover?: (element: T | null) => void;
+}
+
+function getDelay(
+  value: UseElementHoverProps["delay"],
+  prop: "open" | "close",
+  pointerType?: PointerEvent["pointerType"],
+) {
+  if (pointerType && !isMouseLikePointerType(pointerType)) {
+    return 0;
+  }
+  if (typeof value === "number") {
+    return value;
+  }
+  if (typeof value === "function") {
+    const result = value();
+    if (typeof result === "number") {
+      return result;
+    }
+    return result?.[prop];
+  }
+  return value?.[prop];
+}
+
+function getRestMs(value: number | (() => number)) {
+  if (typeof value === "function") {
+    return value();
+  }
+  return value;
+}
+
+/**
+ * Opens the floating element while hovering over elements detected by getElementAtHover.
+ * Similar to CSS `:hover` but with configurable delays and rest behavior.
+ */
+export function useElementHover<T = HTMLElement>(
+  context: FloatingContext,
+  props: UseElementHoverProps<T>,
+): ElementProps {
+  const { open, onOpenChange, dataRef, events, elements } = context;
+  const {
+    attachTo,
+    enabled = true,
+    delay = 0,
+    mouseOnly = false,
+    restMs = 0,
+    getElementAtHover,
+    onHover,
+  } = props;
+
+  const delayRef = useLatestRef(delay);
+  const openRef = useLatestRef(open);
+  const restMsRef = useLatestRef(restMs);
+  const getElementAtHoverRef = useLatestRef(getElementAtHover);
+  const onHoverRef = useLatestRef(onHover);
+
+  const pointerTypeRef = React.useRef<PointerEvent["pointerType"]>(undefined);
+  const timeoutRef = React.useRef(-1);
+  const restTimeoutRef = React.useRef(-1);
+  const blockMouseMoveRef = React.useRef(true);
+  const restTimeoutPendingRef = React.useRef(false);
+  const currentHoveredElementRef = React.useRef<T | null>(null);
+
+  // When closing before opening, clear the delay timeouts to cancel it
+  // from showing.
+  React.useEffect(() => {
+    if (!enabled) {
+      return;
+    }
+
+    function onOpenChange({ open }: { open: boolean }) {
+      if (!open) {
+        clearTimeoutIfSet(timeoutRef);
+        clearTimeoutIfSet(restTimeoutRef);
+        blockMouseMoveRef.current = true;
+        restTimeoutPendingRef.current = false;
+        currentHoveredElementRef.current = null;
+      }
+    }
+
+    events.on("openchange", onOpenChange);
+    return () => {
+      events.off("openchange", onOpenChange);
+    };
+  }, [enabled, events]);
+
+  const closeWithDelay = React.useCallback(
+    (
+      event: Event,
+      runElseBranch = true,
+      reason: OpenChangeReason = "hover",
+    ) => {
+      const closeDelay = getDelay(
+        delayRef.current,
+        "close",
+        pointerTypeRef.current,
+      );
+      if (closeDelay && runElseBranch) {
+        clearTimeoutIfSet(timeoutRef);
+        timeoutRef.current = window.setTimeout(
+          () => onOpenChange(false, event, reason),
+          closeDelay,
+        );
+      } else if (runElseBranch) {
+        clearTimeoutIfSet(timeoutRef);
+        onOpenChange(false, event, reason);
+      }
+    },
+    [delayRef, onOpenChange],
+  );
+
+  const isClickLikeOpenEvent = useEffectEvent(() => {
+    return dataRef.current.openEvent
+      ? ["click", "mousedown"].includes(dataRef.current.openEvent.type)
+      : false;
+  });
+
+  // Registering the mouse events on the attachTo element
+  React.useEffect(() => {
+    if (!enabled) {
+      return;
+    }
+
+    function onMouseEnter(event: MouseEvent) {
+      clearTimeoutIfSet(timeoutRef);
+      blockMouseMoveRef.current = false;
+
+      const element = getElementAtHoverRef.current(event.target as EventTarget);
+
+      // If no element found, don't open
+      if (!element) {
+        currentHoveredElementRef.current = null;
+        onHoverRef.current?.(null);
+        return;
+      }
+
+      // If same element, don't reopen
+      if (currentHoveredElementRef.current === element) {
+        return;
+      }
+
+      currentHoveredElementRef.current = element;
+      onHoverRef.current?.(element);
+
+      if (
+        (mouseOnly && !isMouseLikePointerType(pointerTypeRef.current)) ||
+        (getRestMs(restMsRef.current) > 0 &&
+          !getDelay(delayRef.current, "open"))
+      ) {
+        return;
+      }
+
+      const openDelay = getDelay(
+        delayRef.current,
+        "open",
+        pointerTypeRef.current,
+      );
+
+      if (openDelay) {
+        timeoutRef.current = window.setTimeout(() => {
+          if (!openRef.current) {
+            onOpenChange(true, event, "hover");
+          }
+        }, openDelay);
+      } else if (!open) {
+        onOpenChange(true, event, "hover");
+      }
+    }
+
+    function onMouseLeave(event: MouseEvent) {
+      if (isClickLikeOpenEvent()) {
+        return;
+      }
+
+      clearTimeoutIfSet(restTimeoutRef);
+      restTimeoutPendingRef.current = false;
+
+      // Check if we're leaving the hovered element
+      const element = getElementAtHoverRef.current(event.target as EventTarget);
+      if (element === currentHoveredElementRef.current) {
+        currentHoveredElementRef.current = null;
+        onHoverRef.current?.(null);
+        closeWithDelay(event);
+      }
+    }
+
+    function onMouseMove(event: MouseEvent) {
+      const element = getElementAtHoverRef.current(event.target as EventTarget);
+
+      if (!element) {
+        if (currentHoveredElementRef.current) {
+          currentHoveredElementRef.current = null;
+          onHoverRef.current?.(null);
+          closeWithDelay(event);
+        }
+        return;
+      }
+
+      if (element === currentHoveredElementRef.current) {
+        return;
+      }
+
+      // New element detected
+      currentHoveredElementRef.current = element;
+      onHoverRef.current?.(element);
+
+      if (mouseOnly && !isMouseLikePointerType(pointerTypeRef.current)) {
+        return;
+      }
+
+      if (open || getRestMs(restMsRef.current) === 0) {
+        return;
+      }
+
+      // Ignore insignificant movements to account for tremors.
+      if (
+        restTimeoutPendingRef.current &&
+        (event as MouseEvent).movementX ** 2 +
+          (event as MouseEvent).movementY ** 2 <
+          2
+      ) {
+        return;
+      }
+
+      clearTimeoutIfSet(restTimeoutRef);
+      if (pointerTypeRef.current === "touch") {
+        if (!openRef.current) {
+          onOpenChange(true, event, "hover");
+        }
+      } else {
+        restTimeoutPendingRef.current = true;
+        restTimeoutRef.current = window.setTimeout(() => {
+          if (!openRef.current) {
+            onOpenChange(true, event, "hover");
+          }
+        }, getRestMs(restMsRef.current));
+      }
+    }
+
+    function onFloatingMouseEnter() {
+      clearTimeoutIfSet(timeoutRef);
+    }
+
+    function onFloatingMouseLeave(event: MouseEvent) {
+      if (!isClickLikeOpenEvent()) {
+        closeWithDelay(event, false);
+      }
+    }
+
+    const el = attachTo();
+
+    if (isElement(el)) {
+      const floating = elements.floating;
+
+      el.addEventListener("mouseenter", onMouseEnter);
+      el.addEventListener("mouseleave", onMouseLeave);
+      el.addEventListener("mousemove", onMouseMove);
+
+      if (floating) {
+        floating.addEventListener("mouseenter", onFloatingMouseEnter);
+        floating.addEventListener("mouseleave", onFloatingMouseLeave);
+      }
+
+      return () => {
+        el.removeEventListener("mouseenter", onMouseEnter);
+        el.removeEventListener("mouseleave", onMouseLeave);
+        el.removeEventListener("mousemove", onMouseMove);
+
+        if (floating) {
+          floating.removeEventListener("mouseenter", onFloatingMouseEnter);
+          floating.removeEventListener("mouseleave", onFloatingMouseLeave);
+        }
+      };
+    }
+
+    return;
+  }, [
+    enabled,
+    attachTo,
+    elements,
+    mouseOnly,
+    closeWithDelay,
+    onOpenChange,
+    open,
+    openRef,
+    delayRef,
+    restMsRef,
+    getElementAtHoverRef,
+    onHoverRef,
+    dataRef,
+    isClickLikeOpenEvent,
+  ]);
+
+  useModernLayoutEffect(() => {
+    if (!open) {
+      pointerTypeRef.current = undefined;
+      restTimeoutPendingRef.current = false;
+      currentHoveredElementRef.current = null;
+    }
+  }, [open]);
+
+  React.useEffect(() => {
+    return () => {
+      clearTimeoutIfSet(timeoutRef);
+      clearTimeoutIfSet(restTimeoutRef);
+    };
+  }, [enabled]);
+
+  const reference: ElementProps["reference"] = React.useMemo(() => {
+    function setPointerRef(event: React.PointerEvent) {
+      pointerTypeRef.current = event.pointerType;
+    }
+
+    return {
+      onPointerDown: setPointerRef,
+      onPointerEnter: setPointerRef,
+      onMouseMove(event) {
+        const nativeEvent = "nativeEvent" in event ? event.nativeEvent : event;
+
+        function handleMouseMove() {
+          if (!blockMouseMoveRef.current && !openRef.current) {
+            const element = getElementAtHoverRef.current(
+              nativeEvent.target as EventTarget,
+            );
+            if (element) {
+              currentHoveredElementRef.current = element;
+              onHoverRef.current?.(element);
+              onOpenChange(true, nativeEvent, "hover");
+            }
+          }
+        }
+
+        if (mouseOnly && !isMouseLikePointerType(pointerTypeRef.current)) {
+          return;
+        }
+
+        if (open || getRestMs(restMsRef.current) === 0) {
+          return;
+        }
+
+        // Ignore insignificant movements to account for tremors.
+        if (
+          restTimeoutPendingRef.current &&
+          event.movementX ** 2 + event.movementY ** 2 < 2
+        ) {
+          return;
+        }
+
+        clearTimeoutIfSet(restTimeoutRef);
+        if (pointerTypeRef.current === "touch") {
+          handleMouseMove();
+        } else {
+          restTimeoutPendingRef.current = true;
+          restTimeoutRef.current = window.setTimeout(
+            handleMouseMove,
+            getRestMs(restMsRef.current),
+          );
+        }
+      },
+    };
+  }, [
+    mouseOnly,
+    onOpenChange,
+    open,
+    openRef,
+    restMsRef,
+    getElementAtHoverRef,
+    onHoverRef,
+  ]);
+
+  return React.useMemo(
+    () => (enabled ? { reference } : {}),
+    [enabled, reference],
+  );
+}
diff --git a/packages/react/src/hooks/useExtension.ts b/packages/react/src/hooks/useExtension.ts
new file mode 100644
index 0000000000..182e0cb3f3
--- /dev/null
+++ b/packages/react/src/hooks/useExtension.ts
@@ -0,0 +1,57 @@
+import {
+  BlockNoteEditor,
+  createStore,
+  Extension,
+  ExtensionFactory,
+} from "@blocknote/core";
+import { useStore } from "@tanstack/react-store";
+import { useBlockNoteEditor } from "./useBlockNoteEditor.js";
+
+type Store<T> = ReturnType<typeof createStore<T>>;
+
+/**
+ * Use an extension instance
+ */
+export function useExtension<const T extends ExtensionFactory | Extension>(
+  plugin: T,
+  ctx?: { editor?: BlockNoteEditor<any, any, any> },
+): T extends ExtensionFactory ? NonNullable<ReturnType<T>> : T {
+  // eslint-disable-next-line react-hooks/rules-of-hooks
+  const editor = ctx?.editor ?? useBlockNoteEditor();
+
+  const instance = editor.getExtension(plugin);
+
+  if (!instance) {
+    throw new Error("Extension not found", { cause: { plugin } });
+  }
+
+  return instance as T extends ExtensionFactory
+    ? NonNullable<ReturnType<T>>
+    : T;
+}
+
+type ExtractStore<T> = T extends Store<infer U> ? U : never;
+
+/**
+ * Use the state of an extension
+ */
+export function useExtensionState<
+  T extends ExtensionFactory | Extension,
+  TExtension = T extends ExtensionFactory ? ReturnType<T> : T,
+  TStore = TExtension extends { store: Store<any> }
+    ? TExtension["store"]
+    : never,
+  TSelected = NoInfer<ExtractStore<TStore>>,
+>(
+  plugin: T,
+  ctx?: {
+    editor?: BlockNoteEditor<any, any, any>;
+    selector?: (state: NoInfer<ExtractStore<TStore>>) => TSelected;
+  },
+): TSelected {
+  const { store } = useExtension(plugin, ctx);
+  if (!store) {
+    throw new Error("Store not found on plugin", { cause: { plugin } });
+  }
+  return useStore<ExtractStore<TStore>, TSelected>(store, ctx?.selector as any);
+}
diff --git a/packages/react/src/hooks/useSelectedBlocks.ts b/packages/react/src/hooks/useSelectedBlocks.ts
index 85b4ddfcfd..2f5612aae9 100644
--- a/packages/react/src/hooks/useSelectedBlocks.ts
+++ b/packages/react/src/hooks/useSelectedBlocks.ts
@@ -1,45 +1,19 @@
 import {
-  Block,
   BlockNoteEditor,
   BlockSchema,
   InlineContentSchema,
   StyleSchema,
 } from "@blocknote/core";
-import { useCallback, useState } from "react";
-import { useBlockNoteContext } from "../editor/BlockNoteContext.js";
-import { useEditorContentOrSelectionChange } from "./useEditorContentOrSelectionChange.js";
+import { useEditorState } from "./useEditorState.js";
 
 export function useSelectedBlocks<
   BSchema extends BlockSchema,
   ISchema extends InlineContentSchema,
   SSchema extends StyleSchema,
 >(editor?: BlockNoteEditor<BSchema, ISchema, SSchema>) {
-  const editorContext = useBlockNoteContext<BSchema, ISchema, SSchema>();
-  if (!editor) {
-    editor = editorContext?.editor;
-  }
-
-  if (!editor) {
-    throw new Error(
-      "'editor' is required, either from BlockNoteContext or as a function argument",
-    );
-  }
-
-  const e = editor;
-
-  const [selectedBlocks, setSelectedBlocks] = useState<
-    Block<BSchema, ISchema, SSchema>[]
-  >(() => e.getSelection()?.blocks || [e.getTextCursorPosition().block]);
-
-  const updateSelectedBlocks = useCallback(
-    () =>
-      setSelectedBlocks(
-        e.getSelection()?.blocks || [e.getTextCursorPosition().block],
-      ),
-    [e],
-  );
-
-  useEditorContentOrSelectionChange(updateSelectedBlocks, e);
-
-  return selectedBlocks;
+  return useEditorState({
+    editor,
+    selector: ({ editor }) =>
+      editor.getSelection()?.blocks || [editor.getTextCursorPosition().block],
+  });
 }
diff --git a/packages/react/src/index.ts b/packages/react/src/index.ts
index f1058fb223..28b4f7b197 100644
--- a/packages/react/src/index.ts
+++ b/packages/react/src/index.ts
@@ -57,7 +57,6 @@ export * from "./components/SideMenu/DragHandleMenu/DefaultItems/BlockColorsItem
 export * from "./components/SideMenu/DragHandleMenu/DefaultItems/RemoveBlockItem.js";
 export * from "./components/SideMenu/DragHandleMenu/DefaultItems/TableHeadersItem.js";
 export * from "./components/SideMenu/DragHandleMenu/DragHandleMenu.js";
-export * from "./components/SideMenu/DragHandleMenu/DragHandleMenuProps.js";
 
 export * from "./components/SuggestionMenu/SuggestionMenuController.js";
 export * from "./components/SuggestionMenu/SuggestionMenuWrapper.js";
@@ -90,7 +89,6 @@ export * from "./components/TableHandles/hooks/useTableHandlesPositioning.js";
 export * from "./components/TableHandles/TableCellMenu/DefaultButtons/ColorPicker.js";
 export * from "./components/TableHandles/TableCellMenu/DefaultButtons/SplitButton.js";
 export * from "./components/TableHandles/TableCellMenu/TableCellMenu.js";
-export * from "./components/TableHandles/TableCellMenu/TableCellMenuProps.js";
 export * from "./components/TableHandles/TableCellButton.js";
 export * from "./components/TableHandles/TableCellButtonProps.js";
 
@@ -126,6 +124,7 @@ export * from "./hooks/useSelectedBlocks.js";
 export * from "./hooks/useUIElementPositioning.js";
 export * from "./hooks/useUIPluginState.js";
 export * from "./hooks/useUploadLoading.js";
+export * from "./hooks/useExtension.js";
 
 export * from "./schema/ReactBlockSpec.js";
 export * from "./schema/ReactInlineContentSpec.js";
diff --git a/packages/react/src/schema/ReactBlockSpec.tsx b/packages/react/src/schema/ReactBlockSpec.tsx
index 4abf2a116b..1a341cc341 100644
--- a/packages/react/src/schema/ReactBlockSpec.tsx
+++ b/packages/react/src/schema/ReactBlockSpec.tsx
@@ -3,7 +3,7 @@ import {
   BlockImplementation,
   BlockNoDefaults,
   BlockNoteEditor,
-  BlockNoteExtension,
+  Extension,
   BlockSpec,
   camelToDataKebab,
   CustomBlockImplementation,
@@ -61,7 +61,7 @@ export type ReactCustomBlockSpec<
 > = {
   config: BlockConfig<T, PS, C>;
   implementation: ReactCustomBlockImplementation<T, PS, C>;
-  extensions?: BlockNoteExtension<any>[];
+  extensions?: Extension<any>[];
 };
 
 // Function that wraps the React component returned from 'blockConfig.render' in
@@ -135,10 +135,10 @@ export function createReactBlockSpec<
             options: Partial<TOptions>,
           ) => ReactCustomBlockImplementation<TName, TProps, TContent>),
   extensionsOrCreator?:
-    | BlockNoteExtension<any>[]
+    | Extension<any>[]
     | (TOptions extends undefined
-        ? () => BlockNoteExtension<any>[]
-        : (options: Partial<TOptions>) => BlockNoteExtension<any>[]),
+        ? () => Extension<any>[]
+        : (options: Partial<TOptions>) => Extension<any>[]),
 ): (options?: Partial<TOptions>) => BlockSpec<TName, TProps, TContent>;
 export function createReactBlockSpec<
   const TName extends string,
@@ -168,10 +168,10 @@ export function createReactBlockSpec<
             BlockConf["content"]
           >),
   extensionsOrCreator?:
-    | BlockNoteExtension<any>[]
+    | Extension<any>[]
     | (TOptions extends undefined
-        ? () => BlockNoteExtension<any>[]
-        : (options: Partial<TOptions>) => BlockNoteExtension<any>[]),
+        ? () => Extension<any>[]
+        : (options: Partial<TOptions>) => Extension<any>[]),
 ): (
   options?: Partial<TOptions>,
 ) => BlockSpec<
@@ -198,10 +198,10 @@ export function createReactBlockSpec<
             options: Partial<TOptions>,
           ) => ReactCustomBlockImplementation<TName, TProps, TContent>),
   extensionsOrCreator?:
-    | BlockNoteExtension<any>[]
+    | Extension<any>[]
     | (TOptions extends undefined
-        ? () => BlockNoteExtension<any>[]
-        : (options: Partial<TOptions>) => BlockNoteExtension<any>[]),
+        ? () => Extension<any>[]
+        : (options: Partial<TOptions>) => Extension<any>[]),
 ): (options?: Partial<TOptions>) => BlockSpec<TName, TProps, TContent> {
   return (options = {} as TOptions) => {
     const blockConfig =
diff --git a/packages/xl-ai/src/AIExtension.ts b/packages/xl-ai/src/AIExtension.ts
index 6258cdc38f..5782d966d6 100644
--- a/packages/xl-ai/src/AIExtension.ts
+++ b/packages/xl-ai/src/AIExtension.ts
@@ -1,10 +1,11 @@
 import { Chat } from "@ai-sdk/react";
 import {
-  BlockNoteEditor,
-  BlockNoteExtension,
+  createExtension,
+  createStore,
   getNodeById,
   UnreachableCaseError,
 } from "@blocknote/core";
+import { ForkYDoc, ShowSelection } from "@blocknote/core/extensions";
 import {
   applySuggestions,
   revertSuggestions,
@@ -14,7 +15,6 @@ import { UIMessage } from "ai";
 import { Fragment, Slice } from "prosemirror-model";
 import { Plugin, PluginKey } from "prosemirror-state";
 import { fixTablesKey } from "prosemirror-tables";
-import { createStore, StoreApi } from "zustand/vanilla";
 
 import {
   aiDocumentFormats,
@@ -25,18 +25,7 @@ import {
 import { createAgentCursorPlugin } from "./plugins/AgentCursorPlugin.js";
 import { AIRequestHelpers, InvokeAIOptions } from "./types.js";
 
-type ReadonlyStoreApi<T> = Pick<
-  StoreApi<T>,
-  "getState" | "getInitialState" | "subscribe"
->;
-
 type AIPluginState = {
-  // zustand design considerations:
-  //  - moved this to a nested object to have better typescript typing
-  //  - if we'd do this without a nested object, then we could easily set "wrong" values,
-  //    because "setState" takes a partial object (unless the second parameter "replace" = true),
-  //    and thus we'd lose typescript's typing help
-
   aiMenuState:
     | ({
         /**
@@ -58,67 +47,76 @@ type AIPluginState = {
 
 const PLUGIN_KEY = new PluginKey(`blocknote-ai-plugin`);
 
-export class AIExtension extends BlockNoteExtension {
-  private chatSession:
-    | {
-        previousRequestOptions: InvokeAIOptions;
-        chat: Chat<UIMessage>;
-      }
-    | undefined;
-
-  private scrollInProgress = false;
-  private autoScroll = false;
-
-  public static key(): string {
-    return "ai";
-  }
-
-  // internal store including setters
-  private readonly _store = createStore<AIPluginState>()((_set) => ({
-    aiMenuState: "closed",
-  }));
-
-  /**
-   * Returns a read-only zustand store with the state of the AI Extension
-   */
-  public get store() {
-    // externally, don't expose setters (if we want to do so, expose `set` methods seperately)
-    return this._store as ReadonlyStoreApi<AIPluginState>;
-  }
-
-  /**
-   * Returns a zustand store with the global configuration of the AI Extension.
-   * These options are used by default across all LLM calls when calling {@link executeLLMRequest}
-   */
-  public readonly options: ReturnType<
-    ReturnType<typeof createStore<AIRequestHelpers>>
-  >;
-
-  /**
-   * @internal use `createAIExtension` instead
-   */
-  constructor(
-    public readonly editor: BlockNoteEditor<any, any, any>,
-    options: AIRequestHelpers & {
+declare module "@blocknote/core" {
+  // eslint-disable-next-line @typescript-eslint/no-unused-vars
+  interface BlockNoteEditorOptions<BSchema, ISchema, SSchema> {
+    ai?: AIRequestHelpers & {
       /**
        * The name and color of the agent cursor
        *
        * @default { name: "AI", color: "#8bc6ff" }
        */
       agentCursor?: { name: string; color: string };
-    },
-  ) {
-    super();
+    };
+  }
+}
 
-    this.options = createStore<AIRequestHelpers>()((_set) => ({
-      ...options,
-    }));
+export const AIExtension = createExtension((editor, editorOptions) => {
+  // TODO should we really expose it like this?
+  const options = createStore<AIRequestHelpers>(editorOptions.ai ?? {});
+  const store = createStore<AIPluginState>({
+    aiMenuState: "closed",
+  });
+  let chatSession:
+    | {
+        previousRequestOptions: InvokeAIOptions;
+        chat: Chat<UIMessage>;
+      }
+    | undefined;
+  let autoScroll = false;
+
+  return {
+    key: "ai",
+    options,
+    store,
+    mount({ signal }: { signal: AbortSignal }) {
+      let scrollInProgress = false;
+      // 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 (scrollInProgress) {
+            autoScroll = false;
+          }
 
-    this.addProsemirrorPlugin(
+          scrollInProgress = true;
+        },
+        {
+          capture: true,
+          signal,
+        },
+      );
+      document.addEventListener(
+        "scrollend",
+        () => {
+          scrollInProgress = false;
+        },
+        {
+          capture: true,
+          signal,
+        },
+      );
+    },
+    prosemirrorPlugins: [
       new Plugin({
         key: PLUGIN_KEY,
         filterTransaction: (tr) => {
-          const menuState = this.store.getState().aiMenuState;
+          const menuState = store.state.aiMenuState;
 
           if (menuState !== "closed" && menuState.status === "ai-writing") {
             if (tr.getMeta(fixTablesKey)?.fixTables) {
@@ -131,155 +129,129 @@ export class AIExtension extends BlockNoteExtension {
           return true;
         },
       }),
-    );
-    this.addProsemirrorPlugin(suggestChanges());
-    this.addProsemirrorPlugin(
+      suggestChanges(),
       createAgentCursorPlugin(
-        options.agentCursor || { name: "AI", color: "#8bc6ff" },
+        editorOptions.ai?.agentCursor || { name: "AI", color: "#8bc6ff" },
       ),
-    );
-
-    // 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,
-    );
-  }
+    ],
+
+    /**
+     * Open the AI menu at a specific block
+     */
+    openAIMenuAtBlock(blockID: string) {
+      editor.getExtension(ShowSelection)?.showSelection(true);
+      editor.isEditable = false;
+      store.setState({
+        aiMenuState: {
+          blockId: blockID,
+          status: "user-input",
+        },
+      });
 
-  /**
-   * Open the AI menu at a specific block
-   */
-  public openAIMenuAtBlock(blockID: string) {
-    this.editor.setForceSelectionVisible(true);
-    this.editor.isEditable = false;
-    this._store.setState({
-      aiMenuState: {
-        blockId: blockID,
-        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" });
-  }
+      // Scrolls to the block when the menu opens.
+      const blockElement = editor.domElement?.querySelector(
+        `[data-node-type="blockContainer"][data-id="${blockID}"]`,
+      );
+      blockElement?.scrollIntoView({ block: "center" });
+    },
 
-  /**
-   * Close the AI menu
-   */
-  public closeAIMenu() {
-    this._store.setState({
-      aiMenuState: "closed",
-    });
-    this.chatSession = undefined;
-    this.editor.setForceSelectionVisible(false);
-    this.editor.isEditable = true;
-    this.editor.focus();
-  }
+    /**
+     * Close the AI menu
+     */
+    closeAIMenu() {
+      store.setState({
+        aiMenuState: "closed",
+      });
+      chatSession = undefined;
+      editor.getExtension(ShowSelection)?.showSelection(false);
+      editor.isEditable = true;
+      editor.focus();
+    },
 
-  /**
-   * Accept the changes made by the LLM
-   */
-  public acceptChanges() {
-    // This is slightly convoluted, to try to maintain the undo history as much as possible
-    // The idea is that the LLM call has appended a number of updates to the document, moving the document from state `A` to state `C`
-    // But we want the undo history to skip all of the intermediate states and go straight from `C` to `A`
-    // To do this, we capture the document state `C` (post-LLM call), and then reject the suggestions to recover the original document state `A`
-    // Then we create an intermediate state `B` that captures the diff between `A` and `C`
-    // Then we apply the suggestions to `B` to get the final state `C`
-    // This causes the undo history to skip `B` and go straight from `C` back to `A`
-
-    // Capture the document state `C'` (post-LLM call with all suggestions still in the document)
-    const markedUpDocument = this.editor.prosemirrorState.doc;
-
-    // revert the suggestions to get back to the original document state `A`
-    this.editor.exec((state, dispatch) => {
-      return revertSuggestions(state, (tr) => {
-        dispatch?.(tr.setMeta("addToHistory", false));
+    /**
+     * Accept the changes made by the LLM
+     */
+    acceptChanges() {
+      // This is slightly convoluted, to try to maintain the undo history as much as possible
+      // The idea is that the LLM call has appended a number of updates to the document, moving the document from state `A` to state `C`
+      // But we want the undo history to skip all of the intermediate states and go straight from `C` to `A`
+      // To do this, we capture the document state `C` (post-LLM call), and then reject the suggestions to recover the original document state `A`
+      // Then we create an intermediate state `B` that captures the diff between `A` and `C`
+      // Then we apply the suggestions to `B` to get the final state `C`
+      // This causes the undo history to skip `B` and go straight from `C` back to `A`
+
+      // Capture the document state `C'` (post-LLM call with all suggestions still in the document)
+      const markedUpDocument = editor.prosemirrorState.doc;
+
+      // revert the suggestions to get back to the original document state `A`
+      editor.exec((state, dispatch) => {
+        return revertSuggestions(state, (tr) => {
+          dispatch?.(tr.setMeta("addToHistory", false));
+        });
       });
-    });
-
-    // Create an intermediate state `B` that captures the diff between the original document and the marked up document
-    this.editor.exec((state, dispatch) => {
-      const tr = state.tr;
-      tr.replace(
-        0,
-        tr.doc.content.size,
-        new Slice(Fragment.from(markedUpDocument), 0, 0),
-      );
-      const nextState = state.apply(tr);
-      // Apply the suggestions to the intermediate state `B` to get the final state `C`
-      return applySuggestions(nextState, (resultTr) => {
-        dispatch?.(
-          tr.replace(
-            0,
-            tr.doc.content.size,
-            new Slice(Fragment.from(resultTr.doc), 0, 0),
-          ),
+
+      // Create an intermediate state `B` that captures the diff between the original document and the marked up document
+      editor.exec((state, dispatch) => {
+        const tr = state.tr;
+        tr.replace(
+          0,
+          tr.doc.content.size,
+          new Slice(Fragment.from(markedUpDocument), 0, 0),
         );
+        const nextState = state.apply(tr);
+        // Apply the suggestions to the intermediate state `B` to get the final state `C`
+        return applySuggestions(nextState, (resultTr) => {
+          dispatch?.(
+            tr.replace(
+              0,
+              tr.doc.content.size,
+              new Slice(Fragment.from(resultTr.doc), 0, 0),
+            ),
+          );
+        });
       });
-    });
 
-    // If in collaboration mode, merge the changes back into the original yDoc
-    this.editor.forkYDocPlugin?.merge({ keepChanges: true });
+      // If in collaboration mode, merge the changes back into the original yDoc
+      editor.getExtension(ForkYDoc)?.merge({ keepChanges: true });
 
-    this.closeAIMenu();
-  }
+      this.closeAIMenu();
+    },
 
-  /**
-   * Reject the changes made by the LLM
-   */
-  public rejectChanges() {
-    // Revert the suggestions to get back to the original document
-    this.editor.exec((state, dispatch) => {
-      return revertSuggestions(state, (tr) => {
-        // Do so without adding to history (so the last undo step is just prior to the LLM call)
-        dispatch?.(tr.setMeta("addToHistory", false));
+    /**
+     * Reject the changes made by the LLM
+     */
+    rejectChanges() {
+      // Revert the suggestions to get back to the original document
+      editor.exec((state, dispatch) => {
+        return revertSuggestions(state, (tr) => {
+          // Do so without adding to history (so the last undo step is just prior to the LLM call)
+          dispatch?.(tr.setMeta("addToHistory", false));
+        });
       });
-    });
 
-    // If in collaboration mode, discard the changes and revert to the original yDoc
-    this.editor.forkYDocPlugin?.merge({ keepChanges: false });
-    this.closeAIMenu();
-  }
+      // If in collaboration mode, discard the changes and revert to the original yDoc
+      editor.getExtension(ForkYDoc)?.merge({ keepChanges: false });
+      this.closeAIMenu();
+    },
 
-  /**
-   * Retry the previous LLM call.
-   *
-   * Only valid if the current status is "error"
-   */
-  public async retry() {
-    const { aiMenuState } = this.store.getState();
-    if (
-      aiMenuState === "closed" ||
-      aiMenuState.status !== "error" ||
-      !this.chatSession
-    ) {
-      throw new Error("retry() is only valid when a previous response failed");
-    }
+    /**
+     * Retry the previous LLM call.
+     *
+     * Only valid if the current status is "error"
+     */
+    async retry() {
+      const { aiMenuState } = store.state;
+      if (
+        aiMenuState === "closed" ||
+        aiMenuState.status !== "error" ||
+        !chatSession
+      ) {
+        throw new Error(
+          "retry() is only valid when a previous response failed",
+        );
+      }
 
-    /*
+      /*
     Design decisions:
     - we cannot use chat.regenerate() because the document might have been updated already by toolcalls that failed mid-way
     - we also cannot revert the document changes and use chat.regenerate(), 
@@ -289,198 +261,181 @@ export class AIExtension extends BlockNoteExtension {
       when a call fails
     */
 
-    if (this.chatSession?.chat.status === "error") {
-      // the LLM call failed (i.e. a network error)
-      // console.log("retry failed LLM call", this.chatSession.chat.error);
-      return this.invokeAI({
-        ...this.chatSession.previousRequestOptions,
-        userPrompt: `An error occured in the previous request. Please retry to accomplish the last user prompt.`,
-      });
-    } else {
-      // an error occurred while parsing / executing the previous LLM call
-      // give the LLM a chance to fix the error
-      // console.log("retry failed tool execution");
-      return this.invokeAI({
-        ...this.chatSession.previousRequestOptions,
-        userPrompt: `An error occured while executing the previous tool call. Please retry to accomplish the last user prompt.`,
-      });
-    }
-  }
+      if (chatSession?.chat.status === "error") {
+        // the LLM call failed (i.e. a network error)
+        // console.log("retry failed LLM call", this.chatSession.chat.error);
+        return this.invokeAI({
+          ...chatSession.previousRequestOptions,
+          userPrompt: `An error occured in the previous request. Please retry to accomplish the last user prompt.`,
+        });
+      } else {
+        // an error occurred while parsing / executing the previous LLM call
+        // give the LLM a chance to fix the error
+        // console.log("retry failed tool execution");
+        return this.invokeAI({
+          ...chatSession.previousRequestOptions,
+          userPrompt: `An error occured while executing the previous tool call. Please retry to accomplish the last user prompt.`,
+        });
+      }
+    },
 
-  /**
-   * Update the status of a call to an LLM
-   *
-   * @warning This method should usually only be used for advanced use-cases
-   * if you want to implement how an LLM call is executed. Usually, you should
-   * use {@link executeLLMRequest} instead which will handle the status updates for you.
-   */
-  public setAIResponseStatus(
-    status:
-      | "user-input"
-      | "thinking"
-      | "ai-writing"
-      | "user-reviewing"
-      | {
-          status: "error";
-          error: any;
-        },
-  ) {
-    const aiMenuState = this.store.getState().aiMenuState;
-    if (aiMenuState === "closed") {
-      return;
-    }
-
-    if (status === "ai-writing") {
-      this.editor.setForceSelectionVisible(false);
-    }
-
-    if (typeof status === "object") {
-      if (status.status !== "error") {
-        throw new UnreachableCaseError(status.status);
+    /**
+     * Update the status of a call to an LLM
+     *
+     * @warning This method should usually only be used for advanced use-cases
+     * if you want to implement how an LLM call is executed. Usually, you should
+     * use {@link executeLLMRequest} instead which will handle the status updates for you.
+     */
+    setAIResponseStatus(
+      status:
+        | "user-input"
+        | "thinking"
+        | "ai-writing"
+        | "user-reviewing"
+        | {
+            status: "error";
+            error: any;
+          },
+    ) {
+      const aiMenuState = store.state.aiMenuState;
+      if (aiMenuState === "closed") {
+        return;
       }
-      this._store.setState({
-        aiMenuState: {
-          status: status.status,
-          error: status.error,
-          blockId: aiMenuState.blockId,
-        },
-      });
-    } else {
-      this._store.setState({
-        aiMenuState: {
-          status: status,
-          blockId: aiMenuState.blockId,
-        },
-      });
-    }
-  }
 
-  /**
-   * @deprecated Use {@link invokeAI} instead
-   */
-  public async callLLM(opts: InvokeAIOptions) {
-    return this.invokeAI(opts);
-  }
+      if (status === "ai-writing") {
+        editor.getExtension(ShowSelection)?.showSelection(false);
+      }
 
-  /**
-   * Execute a call to an LLM and apply the result to the editor
-   */
-  public async invokeAI(opts: InvokeAIOptions) {
-    this.setAIResponseStatus("thinking");
-    this.editor.forkYDocPlugin?.fork();
-
-    try {
-      if (!this.chatSession) {
-        // note: in the current implementation opts.transport is only used when creating a new chat
-        // (so changing transport for a subsequent call in the same chat-session is not supported)
-        this.chatSession = {
-          previousRequestOptions: opts,
-          chat: new Chat<UIMessage>({
-            sendAutomaticallyWhen: () => false,
-            transport: opts.transport || this.options.getState().transport,
-          }),
-        };
+      if (typeof status === "object") {
+        if (status.status !== "error") {
+          throw new UnreachableCaseError(status.status);
+        }
+        store.setState({
+          aiMenuState: {
+            status: status.status,
+            error: status.error,
+            blockId: aiMenuState.blockId,
+          },
+        });
       } else {
-        this.chatSession.previousRequestOptions = opts;
+        store.setState({
+          aiMenuState: {
+            status: status,
+            blockId: aiMenuState.blockId,
+          },
+        });
       }
-      const chat = this.chatSession.chat;
-
-      // merge the global options with the local options
-      const globalOpts = this.options.getState();
-      opts = {
-        ...globalOpts,
-        ...opts,
-      } as InvokeAIOptions;
-
-      const sender =
-        opts.aiRequestSender ??
-        defaultAIRequestSender(
-          aiDocumentFormats.html.defaultPromptBuilder,
-          aiDocumentFormats.html.defaultPromptInputDataBuilder,
-        );
-
-      const aiRequest = buildAIRequest({
-        editor: this.editor,
-        chat,
-        userPrompt: opts.userPrompt,
-        useSelection: opts.useSelection,
-        deleteEmptyCursorBlock: opts.deleteEmptyCursorBlock,
-        streamToolsProvider: opts.streamToolsProvider,
-        onBlockUpdated: (blockId) => {
-          // NOTE: does this setState with an anon object trigger unnecessary re-renders?
-          this._store.setState({
-            aiMenuState: {
-              blockId,
-              status: "ai-writing",
-            },
-          });
-
-          // Scrolls to the block being edited by the AI while auto scrolling is
-          // enabled.
-          if (!this.autoScroll) {
-            return;
-          }
-
-          const aiMenuState = this._store.getState().aiMenuState;
-          const aiMenuOpenState =
-            aiMenuState === "closed" ? undefined : aiMenuState;
-          if (!aiMenuOpenState || aiMenuOpenState.status !== "ai-writing") {
-            return;
-          }
+    },
 
-          const nodeInfo = getNodeById(
-            aiMenuOpenState.blockId,
-            this.editor.prosemirrorState.doc,
-          );
-          if (!nodeInfo) {
-            return;
-          }
+    /**
+     * @deprecated Use {@link invokeAI} instead
+     */
+    async callLLM(opts: InvokeAIOptions) {
+      return this.invokeAI(opts);
+    },
 
-          const blockElement = this.editor.prosemirrorView.domAtPos(
-            nodeInfo.posBeforeNode + 1,
+    /**
+     * Execute a call to an LLM and apply the result to the editor
+     */
+    async invokeAI(opts: InvokeAIOptions) {
+      this.setAIResponseStatus("thinking");
+      editor.getExtension(ForkYDoc)?.fork();
+
+      try {
+        if (!chatSession) {
+          // note: in the current implementation opts.transport is only used when creating a new chat
+          // (so changing transport for a subsequent call in the same chat-session is not supported)
+          chatSession = {
+            previousRequestOptions: opts,
+            chat: new Chat<UIMessage>({
+              sendAutomaticallyWhen: () => false,
+              transport: opts.transport || options.state.transport,
+            }),
+          };
+        } else {
+          chatSession.previousRequestOptions = opts;
+        }
+        const chat = chatSession.chat;
+
+        // merge the global options with the local options
+        const globalOpts = options.state;
+        opts = {
+          ...globalOpts,
+          ...opts,
+        } as InvokeAIOptions;
+
+        const sender =
+          opts.aiRequestSender ??
+          defaultAIRequestSender(
+            aiDocumentFormats.html.defaultPromptBuilder,
+            aiDocumentFormats.html.defaultPromptInputDataBuilder,
           );
-          (blockElement.node as HTMLElement).scrollIntoView({
-            block: "center",
-          });
-        },
-      });
 
-      await executeAIRequest({
-        aiRequest,
-        sender,
-        chatRequestOptions: opts.chatRequestOptions,
-        onStart: () => {
-          this.autoScroll = true;
-          this.setAIResponseStatus("ai-writing");
-        },
-      });
+        const aiRequest = buildAIRequest({
+          editor,
+          chat,
+          userPrompt: opts.userPrompt,
+          useSelection: opts.useSelection,
+          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;
+            }
 
-      this.setAIResponseStatus("user-reviewing");
-    } catch (e) {
-      this.setAIResponseStatus({
-        status: "error",
-        error: e,
-      });
-      // eslint-disable-next-line no-console
-      console.warn("Error calling LLM", e, this.chatSession?.chat.messages);
-    }
-  }
-}
+            const aiMenuState = store.state.aiMenuState;
+            const aiMenuOpenState =
+              aiMenuState === "closed" ? undefined : aiMenuState;
+            if (!aiMenuOpenState || aiMenuOpenState.status !== "ai-writing") {
+              return;
+            }
 
-/**
- * Create a new AIExtension instance, this can be passed to the BlockNote editor via the `extensions` option
- */
-export function createAIExtension(
-  options: ConstructorParameters<typeof AIExtension>[1],
-) {
-  return (editor: BlockNoteEditor<any, any, any>) => {
-    return new AIExtension(editor, options);
-  };
-}
+            const nodeInfo = getNodeById(
+              aiMenuOpenState.blockId,
+              editor.prosemirrorState.doc,
+            );
+            if (!nodeInfo) {
+              return;
+            }
 
-/**
- * Return the AIExtension instance from the editor
- */
-export function getAIExtension(editor: BlockNoteEditor<any, any, any>) {
-  return editor.extension(AIExtension);
-}
+            const blockElement = editor.prosemirrorView.domAtPos(
+              nodeInfo.posBeforeNode + 1,
+            );
+            (blockElement.node as HTMLElement).scrollIntoView({
+              block: "center",
+            });
+          },
+        });
+
+        await executeAIRequest({
+          aiRequest,
+          sender,
+          chatRequestOptions: opts.chatRequestOptions,
+          onStart: () => {
+            autoScroll = true;
+            this.setAIResponseStatus("ai-writing");
+          },
+        });
+
+        this.setAIResponseStatus("user-reviewing");
+      } catch (e) {
+        this.setAIResponseStatus({
+          status: "error",
+          error: e,
+        });
+        // eslint-disable-next-line no-console
+        console.warn("Error calling LLM", e, chatSession?.chat.messages);
+      }
+    },
+  } as const;
+});
diff --git a/packages/xl-ai/src/api/formats/json/tools/jsontools.test.ts b/packages/xl-ai/src/api/formats/json/tools/jsontools.test.ts
index 6791375a1f..a1b8ec56e8 100644
--- a/packages/xl-ai/src/api/formats/json/tools/jsontools.test.ts
+++ b/packages/xl-ai/src/api/formats/json/tools/jsontools.test.ts
@@ -13,7 +13,7 @@ import { DeleteBlockToolCall } from "../../base-tools/delete.js";
 import { tools } from "./index.js";
 
 // Helper function to create a mock stream from operations
-import { getAIExtension } from "../../../../AIExtension.js";
+import { AIExtension } from "../../../../AIExtension.js";
 import { StreamToolExecutor } from "../../../../streamTool/StreamToolExecutor.js";
 import { StreamTool } from "../../../../streamTool/streamTool.js";
 import { getExpectedEditor } from "../../../../testUtil/cases/index.js";
@@ -64,7 +64,7 @@ async function executeTestCase(
 
   validateRejectingResultsInOriginalDoc(editor, originalDoc);
 
-  getAIExtension(editor).acceptChanges();
+  editor.getExtension(AIExtension)?.acceptChanges();
   expect(editor.document).toEqual(getExpectedEditor(testCase).document);
 }
 
diff --git a/packages/xl-ai/src/api/formats/tests/sharedTestCases.ts b/packages/xl-ai/src/api/formats/tests/sharedTestCases.ts
index d32533865c..a685aca947 100644
--- a/packages/xl-ai/src/api/formats/tests/sharedTestCases.ts
+++ b/packages/xl-ai/src/api/formats/tests/sharedTestCases.ts
@@ -4,7 +4,7 @@ import { getCurrentTest, TaskContext } from "@vitest/runner";
 import path from "path";
 import { TextSelection } from "prosemirror-state";
 import { describe, expect, it } from "vitest";
-import { getAIExtension } from "../../../AIExtension.js";
+import { AIExtension } from "../../../AIExtension.js";
 import { aiDocumentFormats, defaultAIRequestSender } from "../../../index.js";
 import { addOperationTestCases } from "../../../testUtil/cases/addOperationTestCases.js";
 import { combinedOperationsTestCases } from "../../../testUtil/cases/combinedOperationsTestCases.js";
@@ -118,7 +118,7 @@ export function generateSharedTestCases(
     validateRejectingResultsInOriginalDoc(editor, originalDoc);
 
     // we first need to accept changes to get the correct result
-    getAIExtension(editor).acceptChanges();
+    editor.getExtension(AIExtension)?.acceptChanges();
     expect(editor.document).toEqual(
       getExpectedEditor(test, {
         deleteEmptyCursorBlock: true,
diff --git a/packages/xl-ai/src/components/AIMenu/AIMenu.tsx b/packages/xl-ai/src/components/AIMenu/AIMenu.tsx
index f9e1de4683..41de056987 100644
--- a/packages/xl-ai/src/components/AIMenu/AIMenu.tsx
+++ b/packages/xl-ai/src/components/AIMenu/AIMenu.tsx
@@ -1,10 +1,14 @@
 import { BlockNoteEditor } from "@blocknote/core";
-import { useBlockNoteEditor, useComponentsContext } from "@blocknote/react";
+import {
+  useBlockNoteEditor,
+  useComponentsContext,
+  useExtension,
+  useExtensionState,
+} from "@blocknote/react";
 import { useCallback, useEffect, useMemo, useState } from "react";
 import { RiSparkling2Fill } from "react-icons/ri";
-import { useStore } from "zustand";
 
-import { getAIExtension } from "../../AIExtension.js";
+import { AIExtension } from "../../AIExtension.js";
 import { useAIDictionary } from "../../i18n/useAIDictionary.js";
 import { PromptSuggestionMenu } from "./PromptSuggestionMenu.js";
 import {
@@ -33,11 +37,12 @@ export const AIMenu = (props: AIMenuProps) => {
 
   const Components = useComponentsContext()!;
 
-  const ai = getAIExtension(editor);
+  const ai = useExtension(AIExtension);
 
-  const aiResponseStatus = useStore(ai.store, (state) =>
-    state.aiMenuState !== "closed" ? state.aiMenuState.status : "closed",
-  );
+  const aiResponseStatus = useExtensionState(AIExtension, {
+    selector: (state) =>
+      state.aiMenuState !== "closed" ? state.aiMenuState.status : "closed",
+  });
 
   const { items: externalItems } = props;
   // note, technically there might be a bug with this useMemo when quickly changing the selection and opening the menu
diff --git a/packages/xl-ai/src/components/AIMenu/AIMenuController.tsx b/packages/xl-ai/src/components/AIMenu/AIMenuController.tsx
index fa6d27dce5..947a9c858a 100644
--- a/packages/xl-ai/src/components/AIMenu/AIMenuController.tsx
+++ b/packages/xl-ai/src/components/AIMenu/AIMenuController.tsx
@@ -1,16 +1,19 @@
 import { useBlockNoteEditor } from "@blocknote/react";
 import { FC } from "react";
-import { useStore } from "zustand";
+import { useExtension, useExtensionState } from "@blocknote/react";
 
-import { getAIExtension } from "../../AIExtension.js";
+import { AIExtension } from "../../AIExtension.js";
 import { AIMenu, AIMenuProps } from "./AIMenu.js";
 import { BlockPositioner } from "./BlockPositioner.js";
 
 export const AIMenuController = (props: { aiMenu?: FC<AIMenuProps> }) => {
   const editor = useBlockNoteEditor();
-  const ai = getAIExtension(editor);
+  const ai = useExtension(AIExtension);
 
-  const aiMenuState = useStore(ai.store, (state) => state.aiMenuState);
+  const aiMenuState = useExtensionState(AIExtension, {
+    editor,
+    selector: (state) => state.aiMenuState,
+  });
 
   const blockId = aiMenuState === "closed" ? undefined : aiMenuState.blockId;
 
diff --git a/packages/xl-ai/src/components/AIMenu/PromptSuggestionMenu.tsx b/packages/xl-ai/src/components/AIMenu/PromptSuggestionMenu.tsx
index b18e423a8a..876c11acf9 100644
--- a/packages/xl-ai/src/components/AIMenu/PromptSuggestionMenu.tsx
+++ b/packages/xl-ai/src/components/AIMenu/PromptSuggestionMenu.tsx
@@ -1,4 +1,5 @@
-import { filterSuggestionItems, mergeCSSClasses } from "@blocknote/core";
+import { mergeCSSClasses } from "@blocknote/core";
+import { filterSuggestionItems } from "@blocknote/core/extensions";
 import {
   DefaultReactSuggestionItem,
   useComponentsContext,
diff --git a/packages/xl-ai/src/components/AIMenu/getDefaultAIMenuItems.tsx b/packages/xl-ai/src/components/AIMenu/getDefaultAIMenuItems.tsx
index 33a1c35413..cc5ab779b6 100644
--- a/packages/xl-ai/src/components/AIMenu/getDefaultAIMenuItems.tsx
+++ b/packages/xl-ai/src/components/AIMenu/getDefaultAIMenuItems.tsx
@@ -18,7 +18,7 @@ import {
   RiTextWrap,
 } from "react-icons/ri";
 
-import { getAIExtension } from "../../AIExtension.js";
+import { AIExtension } from "../../AIExtension.js";
 import { aiDocumentFormats } from "../../api/index.js";
 import { getAIDictionary } from "../../i18n/dictionary.js";
 
@@ -40,7 +40,10 @@ function getDefaultAIMenuItemsWithoutSelection<
   S extends StyleSchema,
 >(editor: BlockNoteEditor<BSchema, I, S>): AIMenuSuggestionItem[] {
   const dict = getAIDictionary(editor);
-  const ai = getAIExtension(editor);
+  const ai = editor.getExtension(AIExtension);
+  if (!ai) {
+    return [];
+  }
   return [
     {
       key: "continue_writing",
@@ -128,7 +131,10 @@ function getDefaultAIMenuItemsWithSelection<
 >(editor: BlockNoteEditor<BSchema, I, S>): AIMenuSuggestionItem[] {
   const dict = getAIDictionary(editor);
 
-  const ai = getAIExtension(editor);
+  const ai = editor.getExtension(AIExtension);
+  if (!ai) {
+    return [];
+  }
 
   return [
     {
@@ -216,7 +222,10 @@ function getDefaultAIMenuItemsForReview<
   S extends StyleSchema,
 >(editor: BlockNoteEditor<BSchema, I, S>): AIMenuSuggestionItem[] {
   const dict = getAIDictionary(editor);
-  const ai = getAIExtension(editor);
+  const ai = editor.getExtension(AIExtension);
+  if (!ai) {
+    return [];
+  }
 
   return [
     {
@@ -251,7 +260,10 @@ function getDefaultAIMenuItemsForError<
   S extends StyleSchema,
 >(editor: BlockNoteEditor<BSchema, I, S>): AIMenuSuggestionItem[] {
   const dict = getAIDictionary(editor);
-  const ai = getAIExtension(editor);
+  const ai = editor.getExtension(AIExtension);
+  if (!ai) {
+    return [];
+  }
 
   return [
     {
diff --git a/packages/xl-ai/src/components/FormattingToolbar/AIToolbarButton.tsx b/packages/xl-ai/src/components/FormattingToolbar/AIToolbarButton.tsx
index 1f958fe5c7..6593742a12 100644
--- a/packages/xl-ai/src/components/FormattingToolbar/AIToolbarButton.tsx
+++ b/packages/xl-ai/src/components/FormattingToolbar/AIToolbarButton.tsx
@@ -1,10 +1,11 @@
 import { BlockSchema, InlineContentSchema, StyleSchema } from "@blocknote/core";
-import { useComponentsContext } from "@blocknote/react";
+import { FormattingToolbar } from "@blocknote/core/extensions";
+import { useComponentsContext, useExtension } from "@blocknote/react";
 import { RiSparkling2Fill } from "react-icons/ri";
 
 import { useBlockNoteEditor } from "@blocknote/react";
 
-import { getAIExtension } from "../../AIExtension.js";
+import { AIExtension } from "../../AIExtension.js";
 import { useAIDictionary } from "../../i18n/useAIDictionary.js";
 
 export const AIToolbarButton = () => {
@@ -17,10 +18,10 @@ export const AIToolbarButton = () => {
     StyleSchema
   >();
 
-  const ai = getAIExtension(editor);
+  const ai = useExtension(AIExtension);
 
   const onClick = () => {
-    editor.formattingToolbar.closeMenu();
+    editor.getExtension(FormattingToolbar)?.store.setState(false);
     const selection = editor.getSelection();
     if (!selection) {
       throw new Error("No selection");
diff --git a/packages/xl-ai/src/components/SuggestionMenu/getAISlashMenuItems.tsx b/packages/xl-ai/src/components/SuggestionMenu/getAISlashMenuItems.tsx
index 5599d8a296..107f18aa3a 100644
--- a/packages/xl-ai/src/components/SuggestionMenu/getAISlashMenuItems.tsx
+++ b/packages/xl-ai/src/components/SuggestionMenu/getAISlashMenuItems.tsx
@@ -7,7 +7,7 @@ import {
 import { DefaultReactSuggestionItem } from "@blocknote/react";
 import { RiSparkling2Fill } from "react-icons/ri";
 
-import { getAIExtension } from "../../AIExtension.js";
+import { AIExtension } from "../../AIExtension.js";
 import { getAIDictionary } from "../../i18n/dictionary.js";
 const Icons = {
   AI: RiSparkling2Fill,
@@ -21,7 +21,10 @@ export function getAISlashMenuItems<
   I extends InlineContentSchema,
   S extends StyleSchema,
 >(editor: BlockNoteEditor<BSchema, I, S>): DefaultReactSuggestionItem[] {
-  const ai = getAIExtension(editor);
+  const ai = editor.getExtension(AIExtension);
+  if (!ai) {
+    return [];
+  }
   const items = [
     {
       key: "ai",
diff --git a/packages/xl-ai/src/prosemirror/agent.test.ts b/packages/xl-ai/src/prosemirror/agent.test.ts
index 91cc3a43d5..86f5cd72e0 100644
--- a/packages/xl-ai/src/prosemirror/agent.test.ts
+++ b/packages/xl-ai/src/prosemirror/agent.test.ts
@@ -2,7 +2,7 @@ import { BlockNoteEditor, getBlockInfo, getNodeById } from "@blocknote/core";
 import { Fragment, Slice } from "prosemirror-model";
 import { ReplaceStep, Transform } from "prosemirror-transform";
 import { describe, expect, it } from "vitest";
-import { getAIExtension } from "../AIExtension.js";
+import { AIExtension } from "../AIExtension.js";
 import {
   DocumentOperationTestCase,
   getExpectedEditor,
@@ -247,7 +247,7 @@ async function executeTestCase(
   validateRejectingResultsInOriginalDoc(editor, doc);
   expect(results).toMatchSnapshot();
 
-  getAIExtension(editor).acceptChanges();
+  editor.getExtension(AIExtension)?.acceptChanges();
   expect(editor.document).toEqual(getExpectedEditor(test).document);
 
   return results;
diff --git a/packages/xl-ai/src/testUtil/cases/editors/blockFormatting.ts b/packages/xl-ai/src/testUtil/cases/editors/blockFormatting.ts
index 2d5ae82c33..d8cbf47cce 100644
--- a/packages/xl-ai/src/testUtil/cases/editors/blockFormatting.ts
+++ b/packages/xl-ai/src/testUtil/cases/editors/blockFormatting.ts
@@ -1,5 +1,5 @@
 import { BlockNoteEditor } from "@blocknote/core";
-import { createAIExtension } from "../../../AIExtension.js";
+import { AIExtension } from "../../../AIExtension.js";
 import { schemaWithMention as schema } from "../schemas/mention.js";
 
 export function getEditorWithBlockFormatting() {
@@ -22,7 +22,7 @@ export function getEditorWithBlockFormatting() {
     ],
     trailingBlock: false,
     schema,
-    extensions: [createAIExtension({})],
+    extensions: [AIExtension],
   });
   return editor;
 }
diff --git a/packages/xl-ai/src/testUtil/cases/editors/emptyEditor.ts b/packages/xl-ai/src/testUtil/cases/editors/emptyEditor.ts
index 1d4f963857..4da423096c 100644
--- a/packages/xl-ai/src/testUtil/cases/editors/emptyEditor.ts
+++ b/packages/xl-ai/src/testUtil/cases/editors/emptyEditor.ts
@@ -1,5 +1,5 @@
 import { BlockNoteEditor } from "@blocknote/core";
-import { createAIExtension } from "../../../AIExtension.js";
+import { AIExtension } from "../../../AIExtension.js";
 
 export function getEmptyEditor() {
   const editor = BlockNoteEditor.create({
@@ -10,7 +10,7 @@ export function getEmptyEditor() {
       },
     ],
     trailingBlock: false,
-    extensions: [createAIExtension({})],
+    extensions: [AIExtension],
   });
   return editor;
 }
diff --git a/packages/xl-ai/src/testUtil/cases/editors/formattingAndMentions.ts b/packages/xl-ai/src/testUtil/cases/editors/formattingAndMentions.ts
index 2e4fcd90e9..6213020d5c 100644
--- a/packages/xl-ai/src/testUtil/cases/editors/formattingAndMentions.ts
+++ b/packages/xl-ai/src/testUtil/cases/editors/formattingAndMentions.ts
@@ -1,5 +1,5 @@
 import { BlockNoteEditor } from "@blocknote/core";
-import { createAIExtension } from "../../../AIExtension.js";
+import { AIExtension } from "../../../AIExtension.js";
 import { schemaWithMention as schema } from "../schemas/mention.js";
 
 export function getEditorWithFormattingAndMentions() {
@@ -72,7 +72,7 @@ export function getEditorWithFormattingAndMentions() {
     ],
     trailingBlock: false,
     schema,
-    extensions: [createAIExtension({})],
+    extensions: [AIExtension],
   });
   return editor;
 }
diff --git a/packages/xl-ai/src/testUtil/cases/editors/simpleEditor.ts b/packages/xl-ai/src/testUtil/cases/editors/simpleEditor.ts
index b640e4c8a2..60fcddf797 100644
--- a/packages/xl-ai/src/testUtil/cases/editors/simpleEditor.ts
+++ b/packages/xl-ai/src/testUtil/cases/editors/simpleEditor.ts
@@ -1,5 +1,5 @@
 import { BlockNoteEditor } from "@blocknote/core";
-import { createAIExtension } from "../../../AIExtension.js";
+import { AIExtension } from "../../../AIExtension.js";
 import { schemaWithMention as schema } from "../schemas/mention.js";
 
 export function getSimpleEditor() {
@@ -16,7 +16,7 @@ export function getSimpleEditor() {
     ],
     trailingBlock: false,
     schema,
-    extensions: [createAIExtension({})],
+    extensions: [AIExtension],
   });
   return editor;
 }
@@ -39,7 +39,7 @@ export function getSimpleEditorWithCursorBetweenBlocks() {
     ],
     trailingBlock: false,
     schema,
-    extensions: [createAIExtension({})],
+    extensions: [AIExtension],
   });
   editor.setTextCursorPosition("ref2");
   return editor;
diff --git a/packages/xl-ai/src/testUtil/cases/editors/tables.ts b/packages/xl-ai/src/testUtil/cases/editors/tables.ts
index 704fed1016..f36f538405 100644
--- a/packages/xl-ai/src/testUtil/cases/editors/tables.ts
+++ b/packages/xl-ai/src/testUtil/cases/editors/tables.ts
@@ -1,5 +1,5 @@
 import { BlockNoteEditor } from "@blocknote/core";
-import { createAIExtension } from "../../../AIExtension.js";
+import { AIExtension } from "../../../AIExtension.js";
 import { schemaWithMention as schema } from "../schemas/mention.js";
 
 export function getEditorWithTables() {
@@ -38,7 +38,7 @@ export function getEditorWithTables() {
     ],
     schema,
     trailingBlock: false,
-    extensions: [createAIExtension({})],
+    extensions: [AIExtension],
   });
   return editor;
 }
diff --git a/packages/xl-ai/src/testUtil/cases/updateOperationTestCases.ts b/packages/xl-ai/src/testUtil/cases/updateOperationTestCases.ts
index b196082207..29c71fe4f8 100644
--- a/packages/xl-ai/src/testUtil/cases/updateOperationTestCases.ts
+++ b/packages/xl-ai/src/testUtil/cases/updateOperationTestCases.ts
@@ -1,5 +1,5 @@
 import { BlockNoteEditor, getBlockInfo, getNodeById } from "@blocknote/core";
-import { createAIExtension } from "../../AIExtension.js";
+import { AIExtension } from "../../AIExtension.js";
 import { getEditorWithBlockFormatting } from "./editors/blockFormatting.js";
 import { getEditorWithFormattingAndMentions } from "./editors/formattingAndMentions.js";
 import { DocumentOperationTestCase } from "./index.js";
@@ -683,7 +683,7 @@ export const updateOperationTestCases: DocumentOperationTestCase[] = [
           },
         ],
         schema,
-        extensions: [createAIExtension({})],
+        extensions: [AIExtension],
       });
       return editor;
     },
@@ -736,7 +736,7 @@ export const updateOperationTestCases: DocumentOperationTestCase[] = [
           },
         ],
         schema,
-        extensions: [createAIExtension({})],
+        extensions: [AIExtension],
       });
       return editor;
     },
@@ -771,7 +771,7 @@ export const updateOperationTestCases: DocumentOperationTestCase[] = [
           },
         ],
         schema,
-        extensions: [createAIExtension({})],
+        extensions: [AIExtension],
       });
       return editor;
     },
diff --git a/packages/xl-multi-column/src/extensions/ColumnResize/ColumnResizeExtension.ts b/packages/xl-multi-column/src/extensions/ColumnResize/ColumnResizeExtension.ts
index c9b0c08247..91f6eab122 100644
--- a/packages/xl-multi-column/src/extensions/ColumnResize/ColumnResizeExtension.ts
+++ b/packages/xl-multi-column/src/extensions/ColumnResize/ColumnResizeExtension.ts
@@ -1,4 +1,5 @@
 import { BlockNoteEditor, getNodeById } from "@blocknote/core";
+import { SideMenu } from "@blocknote/core/extensions";
 import { Extension } from "@tiptap/core";
 import { Node } from "prosemirror-model";
 import { Plugin, PluginKey, PluginView } from "prosemirror-state";
@@ -189,7 +190,7 @@ class ColumnResizePluginView implements PluginView {
       this.view.state.tr.setMeta(columnResizePluginKey, newState),
     );
 
-    this.editor.sideMenu.freezeMenu();
+    this.editor.getExtension(SideMenu)?.freezeMenu();
   };
 
   // If the plugin isn't in a resize state, we want to update it to either a
@@ -222,7 +223,10 @@ class ColumnResizePluginView implements PluginView {
 
       // Since the resize bar overlaps the side menu, we don't want to show it
       // if the side menu is already open.
-      if (newState.type === "hover" && this.editor.sideMenu.view?.state?.show) {
+      if (
+        newState.type === "hover" &&
+        this.editor.getExtension(SideMenu)?.store.state?.show
+      ) {
         return;
       }
 
@@ -297,7 +301,7 @@ class ColumnResizePluginView implements PluginView {
       this.view.state.tr.setMeta(columnResizePluginKey, newState),
     );
 
-    this.editor.sideMenu.unfreezeMenu();
+    this.editor.getExtension(SideMenu)?.unfreezeMenu();
   };
 
   destroy() {
diff --git a/packages/xl-multi-column/src/extensions/SuggestionMenu/getMultiColumnSlashMenuItems.tsx b/packages/xl-multi-column/src/extensions/SuggestionMenu/getMultiColumnSlashMenuItems.tsx
index 3be4ce53d9..eb7596a7d7 100644
--- a/packages/xl-multi-column/src/extensions/SuggestionMenu/getMultiColumnSlashMenuItems.tsx
+++ b/packages/xl-multi-column/src/extensions/SuggestionMenu/getMultiColumnSlashMenuItems.tsx
@@ -2,9 +2,9 @@ import {
   BlockNoteEditor,
   BlockSchema,
   InlineContentSchema,
-  insertOrUpdateBlock,
   StyleSchema,
 } from "@blocknote/core";
+import { insertOrUpdateBlockForSlashMenu } from "@blocknote/core/extensions";
 import { DefaultReactSuggestionItem } from "@blocknote/react";
 import { TbColumns2, TbColumns3 } from "react-icons/tb";
 
@@ -40,7 +40,7 @@ export function getMultiColumnSlashMenuItems<
         ...getMultiColumnDictionary(editor).slash_menu.two_columns,
         icon: <TbColumns2 size={18} />,
         onItemClick: () => {
-          insertOrUpdateBlock(editor, {
+          insertOrUpdateBlockForSlashMenu(editor, {
             type: "columnList",
             children: [
               {
@@ -67,7 +67,7 @@ export function getMultiColumnSlashMenuItems<
         ...getMultiColumnDictionary(editor).slash_menu.three_columns,
         icon: <TbColumns3 size={18} />,
         onItemClick: () => {
-          insertOrUpdateBlock(editor, {
+          insertOrUpdateBlockForSlashMenu(editor, {
             type: "columnList",
             children: [
               {
diff --git a/playground/src/style.css b/playground/src/style.css
index ff271345e9..f69602d9f3 100644
--- a/playground/src/style.css
+++ b/playground/src/style.css
@@ -37,6 +37,7 @@ body {
   padding-top: 8px;
   margin: 0 auto;
   max-width: 731px;
+  --bn-ui-base-z-index: 100;
 }
 
 .bn-comment-composer .bn-container {
diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml
index 95618bda35..9fe8091d7c 100644
--- a/pnpm-lock.yaml
+++ b/pnpm-lock.yaml
@@ -132,10 +132,10 @@ importers:
         version: 3.7.1-tiptap3(@types/json-schema@7.0.15)(@types/react-dom@19.2.2(@types/react@19.2.2))(@types/react@19.2.2)(react@19.2.0)
       '@liveblocks/react-blocknote':
         specifier: 3.7.1-tiptap3
-        version: 3.7.1-tiptap3(9de3e2eb190bb0c4d53c80e5b03bc011)
+        version: 3.7.1-tiptap3(a29fa970285ba47e5e1cb6a8a763ba1c)
       '@liveblocks/react-tiptap':
         specifier: 3.7.1-tiptap3
-        version: 3.7.1-tiptap3(98c459f67e13877f32bd9600763a049d)
+        version: 3.7.1-tiptap3(de414c04b2ba68846f97ebf709e8da93)
       '@liveblocks/react-ui':
         specifier: 3.7.1-tiptap3
         version: 3.7.1-tiptap3(@types/json-schema@7.0.15)(@types/react-dom@19.2.2(@types/react@19.2.2))(@types/react@19.2.2)(react-dom@19.2.0(react@19.2.0))(react@19.2.0)
@@ -189,7 +189,7 @@ importers:
         version: 3.13.0
       '@tiptap/core':
         specifier: ^3.0.0
-        version: 3.11.0(@tiptap/pm@3.11.0)
+        version: 3.7.2(@tiptap/pm@3.7.2)
       '@uppy/core':
         specifier: ^3.13.1
         version: 3.13.1
@@ -3426,10 +3426,10 @@ importers:
         version: 3.7.1-tiptap3(@types/json-schema@7.0.15)(@types/react-dom@19.2.2(@types/react@19.2.2))(@types/react@19.2.2)(react@19.2.0)
       '@liveblocks/react-blocknote':
         specifier: 3.7.1-tiptap3
-        version: 3.7.1-tiptap3(9de3e2eb190bb0c4d53c80e5b03bc011)
+        version: 3.7.1-tiptap3(a29fa970285ba47e5e1cb6a8a763ba1c)
       '@liveblocks/react-tiptap':
         specifier: 3.7.1-tiptap3
-        version: 3.7.1-tiptap3(98c459f67e13877f32bd9600763a049d)
+        version: 3.7.1-tiptap3(de414c04b2ba68846f97ebf709e8da93)
       '@liveblocks/react-ui':
         specifier: 3.7.1-tiptap3
         version: 3.7.1-tiptap3(@types/json-schema@7.0.15)(@types/react-dom@19.2.2(@types/react@19.2.2))(@types/react@19.2.2)(react-dom@19.2.0(react@19.2.0))(react@19.2.0)
@@ -3793,7 +3793,7 @@ importers:
         version: 6.0.22(react@19.2.0)
       '@tiptap/core':
         specifier: ^3.0.0
-        version: 3.11.0(@tiptap/pm@3.11.0)
+        version: 3.7.2(@tiptap/pm@3.7.2)
       react:
         specifier: ^19.2.0
         version: 19.2.0
@@ -4452,51 +4452,54 @@ importers:
       '@emoji-mart/data':
         specifier: ^1.2.1
         version: 1.2.1
+      '@handlewithcare/prosemirror-inputrules':
+        specifier: 0.1.3
+        version: 0.1.3(prosemirror-model@1.25.4)(prosemirror-state@1.4.4)(prosemirror-view@1.41.3)
       '@hocuspocus/provider':
         specifier: ^2.15.2
         version: 2.15.3(y-protocols@1.0.6(yjs@13.6.27))(yjs@13.6.27)
       '@shikijs/types':
         specifier: 3.13.0
         version: 3.13.0
+      '@tanstack/store':
+        specifier: 0.7.7
+        version: 0.7.7
       '@tiptap/core':
         specifier: ^3.0.0
-        version: 3.11.0(@tiptap/pm@3.11.0)
+        version: 3.7.2(@tiptap/pm@3.7.2)
       '@tiptap/extension-bold':
         specifier: ^3.7.2
-        version: 3.7.2(@tiptap/core@3.11.0(@tiptap/pm@3.11.0))
+        version: 3.7.2(@tiptap/core@3.7.2(@tiptap/pm@3.7.2))
       '@tiptap/extension-code':
         specifier: ^3.7.2
-        version: 3.7.2(@tiptap/core@3.11.0(@tiptap/pm@3.11.0))
+        version: 3.7.2(@tiptap/core@3.7.2(@tiptap/pm@3.7.2))
       '@tiptap/extension-gapcursor':
         specifier: ^3.7.2
-        version: 3.7.2(@tiptap/extensions@3.4.2(@tiptap/core@3.11.0(@tiptap/pm@3.11.0))(@tiptap/pm@3.11.0))
-      '@tiptap/extension-history':
-        specifier: ^3.7.2
-        version: 3.7.2(@tiptap/extensions@3.4.2(@tiptap/core@3.11.0(@tiptap/pm@3.11.0))(@tiptap/pm@3.11.0))
+        version: 3.7.2(@tiptap/extensions@3.4.2(@tiptap/core@3.7.2(@tiptap/pm@3.7.2))(@tiptap/pm@3.7.2))
       '@tiptap/extension-horizontal-rule':
         specifier: ^3.7.2
-        version: 3.7.2(@tiptap/core@3.11.0(@tiptap/pm@3.11.0))(@tiptap/pm@3.11.0)
+        version: 3.7.2(@tiptap/core@3.7.2(@tiptap/pm@3.7.2))(@tiptap/pm@3.7.2)
       '@tiptap/extension-italic':
         specifier: ^3.7.2
-        version: 3.7.2(@tiptap/core@3.11.0(@tiptap/pm@3.11.0))
+        version: 3.7.2(@tiptap/core@3.7.2(@tiptap/pm@3.7.2))
       '@tiptap/extension-link':
         specifier: ^3.7.2
-        version: 3.7.2(@tiptap/core@3.11.0(@tiptap/pm@3.11.0))(@tiptap/pm@3.11.0)
+        version: 3.7.2(@tiptap/core@3.7.2(@tiptap/pm@3.7.2))(@tiptap/pm@3.7.2)
       '@tiptap/extension-paragraph':
         specifier: ^3.7.2
-        version: 3.7.2(@tiptap/core@3.11.0(@tiptap/pm@3.11.0))
+        version: 3.7.2(@tiptap/core@3.7.2(@tiptap/pm@3.7.2))
       '@tiptap/extension-strike':
         specifier: ^3.7.2
-        version: 3.7.2(@tiptap/core@3.11.0(@tiptap/pm@3.11.0))
+        version: 3.7.2(@tiptap/core@3.7.2(@tiptap/pm@3.7.2))
       '@tiptap/extension-text':
         specifier: ^3.7.2
-        version: 3.7.2(@tiptap/core@3.11.0(@tiptap/pm@3.11.0))
+        version: 3.7.2(@tiptap/core@3.7.2(@tiptap/pm@3.7.2))
       '@tiptap/extension-underline':
         specifier: ^3.7.2
-        version: 3.7.2(@tiptap/core@3.11.0(@tiptap/pm@3.11.0))
+        version: 3.7.2(@tiptap/core@3.7.2(@tiptap/pm@3.7.2))
       '@tiptap/pm':
         specifier: ^3.0.0
-        version: 3.11.0
+        version: 3.7.2
       emoji-mart:
         specifier: ^5.6.0
         version: 5.6.0
@@ -4707,24 +4710,39 @@ importers:
       '@floating-ui/react':
         specifier: ^0.27.16
         version: 0.27.16(react-dom@19.2.0(react@19.2.0))(react@19.2.0)
+      '@floating-ui/utils':
+        specifier: 0.2.10
+        version: 0.2.10
+      '@tanstack/react-store':
+        specifier: 0.7.7
+        version: 0.7.7(react-dom@19.2.0(react@19.2.0))(react@19.2.0)
       '@tiptap/core':
         specifier: ^3.0.0
-        version: 3.11.0(@tiptap/pm@3.11.0)
+        version: 3.7.2(@tiptap/pm@3.7.2)
       '@tiptap/pm':
         specifier: ^3.0.0
-        version: 3.11.0
+        version: 3.7.2
       '@tiptap/react':
-        specifier: ^3.11.0
-        version: 3.11.0(@floating-ui/dom@1.7.4)(@tiptap/core@3.11.0(@tiptap/pm@3.11.0))(@tiptap/pm@3.11.0)(@types/react-dom@19.2.2(@types/react@19.2.2))(@types/react@19.2.2)(react-dom@19.2.0(react@19.2.0))(react@19.2.0)
+        specifier: ^3.7.2
+        version: 3.7.2(@floating-ui/dom@1.7.4)(@tiptap/core@3.7.2(@tiptap/pm@3.7.2))(@tiptap/pm@3.7.2)(@types/react-dom@19.2.2(@types/react@19.2.2))(@types/react@19.2.2)(react-dom@19.2.0(react@19.2.0))(react@19.2.0)
+      '@types/use-sync-external-store':
+        specifier: 1.5.0
+        version: 1.5.0
       emoji-mart:
         specifier: ^5.6.0
         version: 5.6.0
+      fast-deep-equal:
+        specifier: ^3.1.3
+        version: 3.1.3
       lodash.merge:
         specifier: ^4.6.2
         version: 4.6.2
       react-icons:
         specifier: ^5.5.0
         version: 5.5.0(react@19.2.0)
+      use-sync-external-store:
+        specifier: 1.6.0
+        version: 1.6.0(react@19.2.0)
     devDependencies:
       '@types/emoji-mart':
         specifier: ^3.0.14
@@ -4788,10 +4806,10 @@ importers:
         version: link:../react
       '@tiptap/core':
         specifier: ^3.0.0
-        version: 3.11.0(@tiptap/pm@3.11.0)
+        version: 3.7.2(@tiptap/pm@3.7.2)
       '@tiptap/pm':
         specifier: ^3.0.0
-        version: 3.11.0
+        version: 3.7.2
       jsdom:
         specifier: ^25.0.1
         version: 25.0.1(canvas@2.11.2(encoding@0.1.13))
@@ -4967,7 +4985,7 @@ importers:
         version: 0.26.28(react-dom@19.2.0(react@19.2.0))(react@19.2.0)
       '@tiptap/core':
         specifier: ^3.0.0
-        version: 3.11.0(@tiptap/pm@3.11.0)
+        version: 3.7.2(@tiptap/pm@3.7.2)
       ai:
         specifier: ^5.0.76
         version: 5.0.76(zod@4.1.12)
@@ -5286,7 +5304,7 @@ importers:
         version: link:../react
       '@tiptap/core':
         specifier: ^3.0.0
-        version: 3.11.0(@tiptap/pm@3.11.0)
+        version: 3.7.2(@tiptap/pm@3.7.2)
       prosemirror-model:
         specifier: ^1.25.4
         version: 1.25.4
@@ -5532,10 +5550,10 @@ importers:
         version: 3.7.1-tiptap3(@types/json-schema@7.0.15)(@types/react-dom@19.2.2(@types/react@19.2.2))(@types/react@19.2.2)(react@19.2.0)
       '@liveblocks/react-blocknote':
         specifier: 3.7.1-tiptap3
-        version: 3.7.1-tiptap3(9de3e2eb190bb0c4d53c80e5b03bc011)
+        version: 3.7.1-tiptap3(a29fa970285ba47e5e1cb6a8a763ba1c)
       '@liveblocks/react-tiptap':
         specifier: 3.7.1-tiptap3
-        version: 3.7.1-tiptap3(98c459f67e13877f32bd9600763a049d)
+        version: 3.7.1-tiptap3(de414c04b2ba68846f97ebf709e8da93)
       '@liveblocks/react-ui':
         specifier: 3.7.1-tiptap3
         version: 3.7.1-tiptap3(@types/json-schema@7.0.15)(@types/react-dom@19.2.2(@types/react@19.2.2))(@types/react@19.2.2)(react-dom@19.2.0(react@19.2.0))(react@19.2.0)
@@ -5699,7 +5717,7 @@ importers:
         version: 1.51.1
       '@tiptap/pm':
         specifier: ^3.0.0
-        version: 3.11.0
+        version: 3.7.2
       '@types/node':
         specifier: ^20.19.22
         version: 20.19.22
@@ -7153,6 +7171,13 @@ packages:
       fumadocs-core: ^14.0.0 || ^15.0.0
       react: 18.x.x || 19.x.x
 
+  '@handlewithcare/prosemirror-inputrules@0.1.3':
+    resolution: {integrity: sha512-LjGitwgSFHICeU6Mfbt+0Bp4BuWyvHfDYJIf7rq1qdNO88tFcWV3CSqw75o/YbsnUObDgp5Dn+gXIQLRwiyCbg==}
+    peerDependencies:
+      prosemirror-model: ^1.0.0
+      prosemirror-state: ^1.0.0
+      prosemirror-view: ^1.0.0
+
   '@hapi/hoek@9.3.0':
     resolution: {integrity: sha512-/c6rf4UJlmHlC9b5BaNvzAcFv7HZ2QHaV0D4/HNlBdvFnvQq8RI4kYdhyPCl7Xj+oWvTWQ8ujhqS53LIgAe6KQ==}
 
@@ -9750,6 +9775,12 @@ packages:
     peerDependencies:
       vite: ^5.2.0 || ^6 || ^7
 
+  '@tanstack/react-store@0.7.7':
+    resolution: {integrity: sha512-qqT0ufegFRDGSof9D/VqaZgjNgp4tRPHZIJq2+QIHkMUtHjaJ0lYrrXjeIUJvjnTbgPfSD1XgOMEt0lmANn6Zg==}
+    peerDependencies:
+      react: ^16.8.0 || ^17.0.0 || ^18.0.0 || ^19.0.0
+      react-dom: ^16.8.0 || ^17.0.0 || ^18.0.0 || ^19.0.0
+
   '@tanstack/react-table@8.21.3':
     resolution: {integrity: sha512-5nNMTSETP4ykGegmVkhjcS8tTLW6Vl4axfEGQN3v0zdHYbK4UfoqfPChclTrJ4EoK9QynqAu9oUf8VEmrpZ5Ww==}
     engines: {node: '>=12'}
@@ -9763,6 +9794,9 @@ packages:
       react: ^16.8.0 || ^17.0.0 || ^18.0.0 || ^19.0.0
       react-dom: ^16.8.0 || ^17.0.0 || ^18.0.0 || ^19.0.0
 
+  '@tanstack/store@0.7.7':
+    resolution: {integrity: sha512-xa6pTan1bcaqYDS9BDpSiS63qa6EoDkPN9RsRaxHuDdVDNntzq3xNwR5YKTU/V3SkSyC9T4YVOPh2zRQN0nhIQ==}
+
   '@tanstack/table-core@8.21.3':
     resolution: {integrity: sha512-ldZXEhOBb8Is7xLs01fR3YEc3DERiz5silj8tnGkFZytt1abEvl/GhUmCE0PMLaMPTa3Jk4HbKmRlHmu+gCftg==}
     engines: {node: '>=12'}
@@ -9789,8 +9823,8 @@ packages:
       '@types/react-dom':
         optional: true
 
-  '@tiptap/core@3.11.0':
-    resolution: {integrity: sha512-kmS7ZVpHm1EMnW1Wmft9H5ZLM7E0G0NGBx+aGEHGDcNxZBXD2ZUa76CuWjIhOGpwsPbELp684ZdpF2JWoNi4Dg==}
+  '@tiptap/core@3.7.2':
+    resolution: {integrity: sha512-fJwNpTx0aq4UU0HNkxPvPYfNBcTHQ/q5xBUdOB5Mgu6clwGES38jVsNNSudB8g53APUmJIS+2fJbkxl3V+0jww==}
     peerDependencies:
       '@tiptap/pm': ^3.0.0
 
@@ -9799,8 +9833,8 @@ packages:
     peerDependencies:
       '@tiptap/core': ^3.0.0
 
-  '@tiptap/extension-bubble-menu@3.11.0':
-    resolution: {integrity: sha512-P3j9lQ+EZ5Zg/isJzLpCPX7bp7WUBmz8GPs/HPlyMyN2su8LqXntITBZr8IP1JNBlB/wR83k/W0XqdC57mG7cA==}
+  '@tiptap/extension-bubble-menu@3.7.2':
+    resolution: {integrity: sha512-rCJu/X7sZEYWkOwLO342JP06f4giVBECPzr/SzG/fQdAidPW96eilPk3L82w5j24kS9odTlxSLlFlIf6UZ2b9w==}
     peerDependencies:
       '@tiptap/core': ^3.0.0
       '@tiptap/pm': ^3.0.0
@@ -9823,8 +9857,8 @@ packages:
       '@tiptap/pm': ^3.0.0
       y-prosemirror: ^1.2.11
 
-  '@tiptap/extension-floating-menu@3.11.0':
-    resolution: {integrity: sha512-nEHdWZHEJYX1II1oJQ4aeZ8O/Kss4BRbYFXQFGIvPelCfCYEATpUJh3aq3767ARSq40bOWyu+Dcd4SCW0We6Sw==}
+  '@tiptap/extension-floating-menu@3.7.2':
+    resolution: {integrity: sha512-g19ratrXlplYDS29VLQa1y/IM/ro0UFhSS4fQokiQKkazwnA1ZVnebjw8ERYg5lkMm/hiImqstpgdO0LtoivvQ==}
     peerDependencies:
       '@floating-ui/dom': ^1.0.0
       '@tiptap/core': ^3.0.0
@@ -9835,11 +9869,6 @@ packages:
     peerDependencies:
       '@tiptap/extensions': ^3.7.2
 
-  '@tiptap/extension-history@3.7.2':
-    resolution: {integrity: sha512-iAGTUxAr7r+tQ/PtIG94jqTJLy/S/VwE43USfWzXCHvbLn60cPJJG7MTCZxYbd+ZuivZVhEhp3EbzCNmHxjp8A==}
-    peerDependencies:
-      '@tiptap/extensions': ^3.7.2
-
   '@tiptap/extension-horizontal-rule@3.7.2':
     resolution: {integrity: sha512-pN+1hJAVVP3uqtpZ5Rm7z5XUB/NGprK6wExJ04xG117E4rTVcaEb1FnMILY3J3A5XbdC3vHX+cblR8mOl1PAMw==}
     peerDependencies:
@@ -9883,11 +9912,11 @@ packages:
       '@tiptap/core': ^3.0.0
       '@tiptap/pm': ^3.0.0
 
-  '@tiptap/pm@3.11.0':
-    resolution: {integrity: sha512-plCQDLCZIOc92cizB8NNhBRN0szvYR3cx9i5IXo6v9Xsgcun8KHNcJkesc2AyeqdIs0BtOJZaqQ9adHThz8UDw==}
+  '@tiptap/pm@3.7.2':
+    resolution: {integrity: sha512-i2fvXDapwo/TWfHM6STYEbkYyF3qyfN6KEBKPrleX/Z80G5bLxom0gB79TsjLNxTLi6mdf0vTHgAcXMG1avc2g==}
 
-  '@tiptap/react@3.11.0':
-    resolution: {integrity: sha512-SDGei/2DjwmhzsxIQNr6dkB6NxLgXZjQ6hF36NfDm4937r5NLrWrNk5tCsoDQiKZ0DHEzuJ6yZM5C7I7LZLB6w==}
+  '@tiptap/react@3.7.2':
+    resolution: {integrity: sha512-tka4ioSmsGI4TyGZ7jAUoIw8t8DVjr1It0B38vZVLqg8M/ZFgR1NkF50TJ6qAkhy8Uz12AO50so0v79tV2pmEA==}
     peerDependencies:
       '@tiptap/core': ^3.0.0
       '@tiptap/pm': ^3.0.0
@@ -10111,6 +10140,9 @@ packages:
   '@types/use-sync-external-store@0.0.6':
     resolution: {integrity: sha512-zFDAD+tlpf2r4asuHEj0XH6pY6i0g5NeAHPn+15wk3BV6JA69eERFXC1gyGThDkVa1zCyKr5jox1+2LbV/AMLg==}
 
+  '@types/use-sync-external-store@1.5.0':
+    resolution: {integrity: sha512-5dyB8nLC/qogMrlCizZnYWQTA4lnb/v+It+sqNl5YnSRAPMlIqY/X0Xn+gZw8vOL+TgTTr28VEbn3uf8fUtAkw==}
+
   '@types/uuid@8.3.4':
     resolution: {integrity: sha512-c/I8ZRb51j+pYGAu5CrFMRxqZ2ke4y2grEBO5AUjgSkSk+qT2Ea+OdWElz/OiMf5MNpn2b17kuVBwZLQJXzihw==}
 
@@ -17630,6 +17662,14 @@ snapshots:
     transitivePeerDependencies:
       - supports-color
 
+  '@handlewithcare/prosemirror-inputrules@0.1.3(prosemirror-model@1.25.4)(prosemirror-state@1.4.4)(prosemirror-view@1.41.3)':
+    dependencies:
+      prosemirror-history: 1.4.1
+      prosemirror-model: 1.25.4
+      prosemirror-state: 1.4.4
+      prosemirror-transform: 1.10.5
+      prosemirror-view: 1.41.3
+
   '@hapi/hoek@9.3.0': {}
 
   '@hapi/topo@5.1.0':
@@ -17919,17 +17959,17 @@ snapshots:
     dependencies:
       '@types/json-schema': 7.0.15
 
-  '@liveblocks/react-blocknote@3.7.1-tiptap3(9de3e2eb190bb0c4d53c80e5b03bc011)':
+  '@liveblocks/react-blocknote@3.7.1-tiptap3(a29fa970285ba47e5e1cb6a8a763ba1c)':
     dependencies:
       '@blocknote/core': link:packages/core
       '@blocknote/react': link:packages/react
       '@liveblocks/client': 3.7.1-tiptap3(@types/json-schema@7.0.15)
       '@liveblocks/core': 3.7.1-tiptap3(@types/json-schema@7.0.15)
       '@liveblocks/react': 3.7.1-tiptap3(@types/json-schema@7.0.15)(@types/react-dom@19.2.2(@types/react@19.2.2))(@types/react@19.2.2)(react@19.2.0)
-      '@liveblocks/react-tiptap': 3.7.1-tiptap3(98c459f67e13877f32bd9600763a049d)
+      '@liveblocks/react-tiptap': 3.7.1-tiptap3(de414c04b2ba68846f97ebf709e8da93)
       '@liveblocks/react-ui': 3.7.1-tiptap3(@types/json-schema@7.0.15)(@types/react-dom@19.2.2(@types/react@19.2.2))(@types/react@19.2.2)(react-dom@19.2.0(react@19.2.0))(react@19.2.0)
       '@liveblocks/yjs': 3.7.1-tiptap3(@types/json-schema@7.0.15)(yjs@13.6.27)
-      '@tiptap/core': 3.11.0(@tiptap/pm@3.11.0)
+      '@tiptap/core': 3.7.2(@tiptap/pm@3.7.2)
       react: 19.2.0
       react-dom: 19.2.0(react@19.2.0)
       vitest-tsconfig-paths: 3.4.1
@@ -17950,7 +17990,7 @@ snapshots:
       - y-protocols
       - yjs
 
-  '@liveblocks/react-tiptap@3.7.1-tiptap3(98c459f67e13877f32bd9600763a049d)':
+  '@liveblocks/react-tiptap@3.7.1-tiptap3(de414c04b2ba68846f97ebf709e8da93)':
     dependencies:
       '@floating-ui/react-dom': 2.1.6(react-dom@19.2.0(react@19.2.0))(react@19.2.0)
       '@liveblocks/client': 3.7.1-tiptap3(@types/json-schema@7.0.15)
@@ -17960,12 +18000,12 @@ snapshots:
       '@liveblocks/yjs': 3.7.1-tiptap3(@types/json-schema@7.0.15)(yjs@13.6.27)
       '@radix-ui/react-select': 2.2.6(@types/react-dom@19.2.2(@types/react@19.2.2))(@types/react@19.2.2)(react-dom@19.2.0(react@19.2.0))(react@19.2.0)
       '@radix-ui/react-toggle': 1.1.10(@types/react-dom@19.2.2(@types/react@19.2.2))(@types/react@19.2.2)(react-dom@19.2.0(react@19.2.0))(react@19.2.0)
-      '@tiptap/core': 3.11.0(@tiptap/pm@3.11.0)
-      '@tiptap/extension-collaboration': 2.26.1(@tiptap/core@3.11.0(@tiptap/pm@3.11.0))(@tiptap/pm@3.11.0)(y-prosemirror@1.3.7(prosemirror-model@1.25.4)(prosemirror-state@1.4.4)(prosemirror-view@1.41.3)(y-protocols@1.0.6(yjs@13.6.27))(yjs@13.6.27))
-      '@tiptap/extension-collaboration-cursor': 2.26.1(@tiptap/core@3.11.0(@tiptap/pm@3.11.0))(y-prosemirror@1.3.7(prosemirror-model@1.25.4)(prosemirror-state@1.4.4)(prosemirror-view@1.41.3)(y-protocols@1.0.6(yjs@13.6.27))(yjs@13.6.27))
-      '@tiptap/pm': 3.11.0
-      '@tiptap/react': 3.11.0(@floating-ui/dom@1.7.4)(@tiptap/core@3.11.0(@tiptap/pm@3.11.0))(@tiptap/pm@3.11.0)(@types/react-dom@19.2.2(@types/react@19.2.2))(@types/react@19.2.2)(react-dom@19.2.0(react@19.2.0))(react@19.2.0)
-      '@tiptap/suggestion': 2.26.1(@tiptap/core@3.11.0(@tiptap/pm@3.11.0))(@tiptap/pm@3.11.0)
+      '@tiptap/core': 3.7.2(@tiptap/pm@3.7.2)
+      '@tiptap/extension-collaboration': 2.26.1(@tiptap/core@3.7.2(@tiptap/pm@3.7.2))(@tiptap/pm@3.7.2)(y-prosemirror@1.3.7(prosemirror-model@1.25.4)(prosemirror-state@1.4.4)(prosemirror-view@1.41.3)(y-protocols@1.0.6(yjs@13.6.27))(yjs@13.6.27))
+      '@tiptap/extension-collaboration-cursor': 2.26.1(@tiptap/core@3.7.2(@tiptap/pm@3.7.2))(y-prosemirror@1.3.7(prosemirror-model@1.25.4)(prosemirror-state@1.4.4)(prosemirror-view@1.41.3)(y-protocols@1.0.6(yjs@13.6.27))(yjs@13.6.27))
+      '@tiptap/pm': 3.7.2
+      '@tiptap/react': 3.7.2(@floating-ui/dom@1.7.4)(@tiptap/core@3.7.2(@tiptap/pm@3.7.2))(@tiptap/pm@3.7.2)(@types/react-dom@19.2.2(@types/react@19.2.2))(@types/react@19.2.2)(react-dom@19.2.0(react@19.2.0))(react@19.2.0)
+      '@tiptap/suggestion': 2.26.1(@tiptap/core@3.7.2(@tiptap/pm@3.7.2))(@tiptap/pm@3.7.2)
       cmdk: 1.1.1(@types/react-dom@19.2.2(@types/react@19.2.2))(@types/react@19.2.2)(react-dom@19.2.0(react@19.2.0))(react@19.2.0)
       react: 19.2.0
       react-dom: 19.2.0(react@19.2.0)
@@ -20705,6 +20745,13 @@ snapshots:
       tailwindcss: 4.1.14
       vite: 6.4.0(@types/node@24.8.1)(jiti@2.6.1)(lightningcss@1.30.1)(terser@5.44.1)(tsx@4.20.6)(yaml@2.8.1)
 
+  '@tanstack/react-store@0.7.7(react-dom@19.2.0(react@19.2.0))(react@19.2.0)':
+    dependencies:
+      '@tanstack/store': 0.7.7
+      react: 19.2.0
+      react-dom: 19.2.0(react@19.2.0)
+      use-sync-external-store: 1.6.0(react@19.2.0)
+
   '@tanstack/react-table@8.21.3(react-dom@19.2.0(react@19.2.0))(react@19.2.0)':
     dependencies:
       '@tanstack/table-core': 8.21.3
@@ -20717,6 +20764,8 @@ snapshots:
       react: 19.2.0
       react-dom: 19.2.0(react@19.2.0)
 
+  '@tanstack/store@0.7.7': {}
+
   '@tanstack/table-core@8.21.3': {}
 
   '@tanstack/virtual-core@3.13.12': {}
@@ -20742,88 +20791,84 @@ snapshots:
       '@types/react': 19.2.2
       '@types/react-dom': 19.2.2(@types/react@19.2.2)
 
-  '@tiptap/core@3.11.0(@tiptap/pm@3.11.0)':
+  '@tiptap/core@3.7.2(@tiptap/pm@3.7.2)':
     dependencies:
-      '@tiptap/pm': 3.11.0
+      '@tiptap/pm': 3.7.2
 
-  '@tiptap/extension-bold@3.7.2(@tiptap/core@3.11.0(@tiptap/pm@3.11.0))':
+  '@tiptap/extension-bold@3.7.2(@tiptap/core@3.7.2(@tiptap/pm@3.7.2))':
     dependencies:
-      '@tiptap/core': 3.11.0(@tiptap/pm@3.11.0)
+      '@tiptap/core': 3.7.2(@tiptap/pm@3.7.2)
 
-  '@tiptap/extension-bubble-menu@3.11.0(@tiptap/core@3.11.0(@tiptap/pm@3.11.0))(@tiptap/pm@3.11.0)':
+  '@tiptap/extension-bubble-menu@3.7.2(@tiptap/core@3.7.2(@tiptap/pm@3.7.2))(@tiptap/pm@3.7.2)':
     dependencies:
       '@floating-ui/dom': 1.7.4
-      '@tiptap/core': 3.11.0(@tiptap/pm@3.11.0)
-      '@tiptap/pm': 3.11.0
+      '@tiptap/core': 3.7.2(@tiptap/pm@3.7.2)
+      '@tiptap/pm': 3.7.2
     optional: true
 
-  '@tiptap/extension-code@3.7.2(@tiptap/core@3.11.0(@tiptap/pm@3.11.0))':
+  '@tiptap/extension-code@3.7.2(@tiptap/core@3.7.2(@tiptap/pm@3.7.2))':
     dependencies:
-      '@tiptap/core': 3.11.0(@tiptap/pm@3.11.0)
+      '@tiptap/core': 3.7.2(@tiptap/pm@3.7.2)
 
-  '@tiptap/extension-collaboration-cursor@2.26.1(@tiptap/core@3.11.0(@tiptap/pm@3.11.0))(y-prosemirror@1.3.7(prosemirror-model@1.25.4)(prosemirror-state@1.4.4)(prosemirror-view@1.41.3)(y-protocols@1.0.6(yjs@13.6.27))(yjs@13.6.27))':
+  '@tiptap/extension-collaboration-cursor@2.26.1(@tiptap/core@3.7.2(@tiptap/pm@3.7.2))(y-prosemirror@1.3.7(prosemirror-model@1.25.4)(prosemirror-state@1.4.4)(prosemirror-view@1.41.3)(y-protocols@1.0.6(yjs@13.6.27))(yjs@13.6.27))':
     dependencies:
-      '@tiptap/core': 3.11.0(@tiptap/pm@3.11.0)
+      '@tiptap/core': 3.7.2(@tiptap/pm@3.7.2)
       y-prosemirror: 1.3.7(prosemirror-model@1.25.4)(prosemirror-state@1.4.4)(prosemirror-view@1.41.3)(y-protocols@1.0.6(yjs@13.6.27))(yjs@13.6.27)
 
-  '@tiptap/extension-collaboration@2.26.1(@tiptap/core@3.11.0(@tiptap/pm@3.11.0))(@tiptap/pm@3.11.0)(y-prosemirror@1.3.7(prosemirror-model@1.25.4)(prosemirror-state@1.4.4)(prosemirror-view@1.41.3)(y-protocols@1.0.6(yjs@13.6.27))(yjs@13.6.27))':
+  '@tiptap/extension-collaboration@2.26.1(@tiptap/core@3.7.2(@tiptap/pm@3.7.2))(@tiptap/pm@3.7.2)(y-prosemirror@1.3.7(prosemirror-model@1.25.4)(prosemirror-state@1.4.4)(prosemirror-view@1.41.3)(y-protocols@1.0.6(yjs@13.6.27))(yjs@13.6.27))':
     dependencies:
-      '@tiptap/core': 3.11.0(@tiptap/pm@3.11.0)
-      '@tiptap/pm': 3.11.0
+      '@tiptap/core': 3.7.2(@tiptap/pm@3.7.2)
+      '@tiptap/pm': 3.7.2
       y-prosemirror: 1.3.7(prosemirror-model@1.25.4)(prosemirror-state@1.4.4)(prosemirror-view@1.41.3)(y-protocols@1.0.6(yjs@13.6.27))(yjs@13.6.27)
 
-  '@tiptap/extension-floating-menu@3.11.0(@floating-ui/dom@1.7.4)(@tiptap/core@3.11.0(@tiptap/pm@3.11.0))(@tiptap/pm@3.11.0)':
+  '@tiptap/extension-floating-menu@3.7.2(@floating-ui/dom@1.7.4)(@tiptap/core@3.7.2(@tiptap/pm@3.7.2))(@tiptap/pm@3.7.2)':
     dependencies:
       '@floating-ui/dom': 1.7.4
-      '@tiptap/core': 3.11.0(@tiptap/pm@3.11.0)
-      '@tiptap/pm': 3.11.0
+      '@tiptap/core': 3.7.2(@tiptap/pm@3.7.2)
+      '@tiptap/pm': 3.7.2
     optional: true
 
-  '@tiptap/extension-gapcursor@3.7.2(@tiptap/extensions@3.4.2(@tiptap/core@3.11.0(@tiptap/pm@3.11.0))(@tiptap/pm@3.11.0))':
-    dependencies:
-      '@tiptap/extensions': 3.4.2(@tiptap/core@3.11.0(@tiptap/pm@3.11.0))(@tiptap/pm@3.11.0)
-
-  '@tiptap/extension-history@3.7.2(@tiptap/extensions@3.4.2(@tiptap/core@3.11.0(@tiptap/pm@3.11.0))(@tiptap/pm@3.11.0))':
+  '@tiptap/extension-gapcursor@3.7.2(@tiptap/extensions@3.4.2(@tiptap/core@3.7.2(@tiptap/pm@3.7.2))(@tiptap/pm@3.7.2))':
     dependencies:
-      '@tiptap/extensions': 3.4.2(@tiptap/core@3.11.0(@tiptap/pm@3.11.0))(@tiptap/pm@3.11.0)
+      '@tiptap/extensions': 3.4.2(@tiptap/core@3.7.2(@tiptap/pm@3.7.2))(@tiptap/pm@3.7.2)
 
-  '@tiptap/extension-horizontal-rule@3.7.2(@tiptap/core@3.11.0(@tiptap/pm@3.11.0))(@tiptap/pm@3.11.0)':
+  '@tiptap/extension-horizontal-rule@3.7.2(@tiptap/core@3.7.2(@tiptap/pm@3.7.2))(@tiptap/pm@3.7.2)':
     dependencies:
-      '@tiptap/core': 3.11.0(@tiptap/pm@3.11.0)
-      '@tiptap/pm': 3.11.0
+      '@tiptap/core': 3.7.2(@tiptap/pm@3.7.2)
+      '@tiptap/pm': 3.7.2
 
-  '@tiptap/extension-italic@3.7.2(@tiptap/core@3.11.0(@tiptap/pm@3.11.0))':
+  '@tiptap/extension-italic@3.7.2(@tiptap/core@3.7.2(@tiptap/pm@3.7.2))':
     dependencies:
-      '@tiptap/core': 3.11.0(@tiptap/pm@3.11.0)
+      '@tiptap/core': 3.7.2(@tiptap/pm@3.7.2)
 
-  '@tiptap/extension-link@3.7.2(@tiptap/core@3.11.0(@tiptap/pm@3.11.0))(@tiptap/pm@3.11.0)':
+  '@tiptap/extension-link@3.7.2(@tiptap/core@3.7.2(@tiptap/pm@3.7.2))(@tiptap/pm@3.7.2)':
     dependencies:
-      '@tiptap/core': 3.11.0(@tiptap/pm@3.11.0)
-      '@tiptap/pm': 3.11.0
+      '@tiptap/core': 3.7.2(@tiptap/pm@3.7.2)
+      '@tiptap/pm': 3.7.2
       linkifyjs: 4.3.2
 
-  '@tiptap/extension-paragraph@3.7.2(@tiptap/core@3.11.0(@tiptap/pm@3.11.0))':
+  '@tiptap/extension-paragraph@3.7.2(@tiptap/core@3.7.2(@tiptap/pm@3.7.2))':
     dependencies:
-      '@tiptap/core': 3.11.0(@tiptap/pm@3.11.0)
+      '@tiptap/core': 3.7.2(@tiptap/pm@3.7.2)
 
-  '@tiptap/extension-strike@3.7.2(@tiptap/core@3.11.0(@tiptap/pm@3.11.0))':
+  '@tiptap/extension-strike@3.7.2(@tiptap/core@3.7.2(@tiptap/pm@3.7.2))':
     dependencies:
-      '@tiptap/core': 3.11.0(@tiptap/pm@3.11.0)
+      '@tiptap/core': 3.7.2(@tiptap/pm@3.7.2)
 
-  '@tiptap/extension-text@3.7.2(@tiptap/core@3.11.0(@tiptap/pm@3.11.0))':
+  '@tiptap/extension-text@3.7.2(@tiptap/core@3.7.2(@tiptap/pm@3.7.2))':
     dependencies:
-      '@tiptap/core': 3.11.0(@tiptap/pm@3.11.0)
+      '@tiptap/core': 3.7.2(@tiptap/pm@3.7.2)
 
-  '@tiptap/extension-underline@3.7.2(@tiptap/core@3.11.0(@tiptap/pm@3.11.0))':
+  '@tiptap/extension-underline@3.7.2(@tiptap/core@3.7.2(@tiptap/pm@3.7.2))':
     dependencies:
-      '@tiptap/core': 3.11.0(@tiptap/pm@3.11.0)
+      '@tiptap/core': 3.7.2(@tiptap/pm@3.7.2)
 
-  '@tiptap/extensions@3.4.2(@tiptap/core@3.11.0(@tiptap/pm@3.11.0))(@tiptap/pm@3.11.0)':
+  '@tiptap/extensions@3.4.2(@tiptap/core@3.7.2(@tiptap/pm@3.7.2))(@tiptap/pm@3.7.2)':
     dependencies:
-      '@tiptap/core': 3.11.0(@tiptap/pm@3.11.0)
-      '@tiptap/pm': 3.11.0
+      '@tiptap/core': 3.7.2(@tiptap/pm@3.7.2)
+      '@tiptap/pm': 3.7.2
 
-  '@tiptap/pm@3.11.0':
+  '@tiptap/pm@3.7.2':
     dependencies:
       prosemirror-changeset: 2.3.1
       prosemirror-collab: 1.3.1
@@ -20844,10 +20889,10 @@ snapshots:
       prosemirror-transform: 1.10.5
       prosemirror-view: 1.41.3
 
-  '@tiptap/react@3.11.0(@floating-ui/dom@1.7.4)(@tiptap/core@3.11.0(@tiptap/pm@3.11.0))(@tiptap/pm@3.11.0)(@types/react-dom@19.2.2(@types/react@19.2.2))(@types/react@19.2.2)(react-dom@19.2.0(react@19.2.0))(react@19.2.0)':
+  '@tiptap/react@3.7.2(@floating-ui/dom@1.7.4)(@tiptap/core@3.7.2(@tiptap/pm@3.7.2))(@tiptap/pm@3.7.2)(@types/react-dom@19.2.2(@types/react@19.2.2))(@types/react@19.2.2)(react-dom@19.2.0(react@19.2.0))(react@19.2.0)':
     dependencies:
-      '@tiptap/core': 3.11.0(@tiptap/pm@3.11.0)
-      '@tiptap/pm': 3.11.0
+      '@tiptap/core': 3.7.2(@tiptap/pm@3.7.2)
+      '@tiptap/pm': 3.7.2
       '@types/react': 19.2.2
       '@types/react-dom': 19.2.2(@types/react@19.2.2)
       '@types/use-sync-external-store': 0.0.6
@@ -20856,15 +20901,15 @@ snapshots:
       react-dom: 19.2.0(react@19.2.0)
       use-sync-external-store: 1.6.0(react@19.2.0)
     optionalDependencies:
-      '@tiptap/extension-bubble-menu': 3.11.0(@tiptap/core@3.11.0(@tiptap/pm@3.11.0))(@tiptap/pm@3.11.0)
-      '@tiptap/extension-floating-menu': 3.11.0(@floating-ui/dom@1.7.4)(@tiptap/core@3.11.0(@tiptap/pm@3.11.0))(@tiptap/pm@3.11.0)
+      '@tiptap/extension-bubble-menu': 3.7.2(@tiptap/core@3.7.2(@tiptap/pm@3.7.2))(@tiptap/pm@3.7.2)
+      '@tiptap/extension-floating-menu': 3.7.2(@floating-ui/dom@1.7.4)(@tiptap/core@3.7.2(@tiptap/pm@3.7.2))(@tiptap/pm@3.7.2)
     transitivePeerDependencies:
       - '@floating-ui/dom'
 
-  '@tiptap/suggestion@2.26.1(@tiptap/core@3.11.0(@tiptap/pm@3.11.0))(@tiptap/pm@3.11.0)':
+  '@tiptap/suggestion@2.26.1(@tiptap/core@3.7.2(@tiptap/pm@3.7.2))(@tiptap/pm@3.7.2)':
     dependencies:
-      '@tiptap/core': 3.11.0(@tiptap/pm@3.11.0)
-      '@tiptap/pm': 3.11.0
+      '@tiptap/core': 3.7.2(@tiptap/pm@3.7.2)
+      '@tiptap/pm': 3.7.2
 
   '@tootallnate/once@2.0.0': {}
 
@@ -21107,6 +21152,8 @@ snapshots:
 
   '@types/use-sync-external-store@0.0.6': {}
 
+  '@types/use-sync-external-store@1.5.0': {}
+
   '@types/uuid@8.3.4': {}
 
   '@types/yargs-parser@21.0.3': {}
diff --git a/tests/src/utils/components/Editor.tsx b/tests/src/utils/components/Editor.tsx
index 3b55e5200b..459920ed2a 100644
--- a/tests/src/utils/components/Editor.tsx
+++ b/tests/src/utils/components/Editor.tsx
@@ -1,4 +1,5 @@
-import { BlockNoteSchema, filterSuggestionItems } from "@blocknote/core";
+import { BlockNoteSchema } from "@blocknote/core";
+import { filterSuggestionItems } from "@blocknote/core/extensions";
 import "@blocknote/core/style.css";
 import { BlockNoteView } from "@blocknote/mantine";
 import "@blocknote/mantine/style.css";
diff --git a/tests/tsconfig.json b/tests/tsconfig.json
index eed3ac7d53..e3ce1824b0 100644
--- a/tests/tsconfig.json
+++ b/tests/tsconfig.json
@@ -4,7 +4,7 @@
     "useDefineForClassFields": true,
     "module": "ESNext",
     "lib": ["ESNext", "DOM"],
-    "moduleResolution": "Node",
+    "moduleResolution": "bundler",
     "jsx": "react-jsx",
     "strict": true,
     "sourceMap": true,