diff --git a/packages/react/src/components/AttributionTooltip/AttributionTooltipController.tsx b/packages/react/src/components/AttributionTooltip/AttributionTooltipController.tsx
index 57a691638e..5d070c0c99 100644
--- a/packages/react/src/components/AttributionTooltip/AttributionTooltipController.tsx
+++ b/packages/react/src/components/AttributionTooltip/AttributionTooltipController.tsx
@@ -1,8 +1,4 @@
-import {
-  AttributionExtension,
-  getReferenceClientRects,
-  getReferenceRect,
-} from "@blocknote/core/y";
+import type { AttributionExtension } from "@blocknote/core/y";
 import { flip, offset, shift, inline } from "@floating-ui/react";
 import { FC, useMemo } from "react";
 
@@ -39,22 +35,33 @@ export const AttributionTooltipController = (props: {
    */
   portalElement?: HTMLElement | null;
 }) => {
-  const state = useExtensionState(AttributionExtension, {
+  const state = useExtensionState<typeof AttributionExtension>("attribution", {
     selector: (state) => state,
   });
 
-  // Position off the hovered mark. Block-level marks are `display: contents`, so
-  // a live `getBoundingClientRect` (via `getReferenceRect`) is provided rather
-  // than relying on the wrapper's own (zero-sized) box.
   const reference = useMemo<GenericPopoverReference | undefined>(
     () =>
       state
         ? {
             element: state.anchor,
-            getBoundingClientRect: () => getReferenceRect(state.anchor),
-            // Required by the `inline()` middleware — a virtual reference has no
-            // default `getClientRects`, and `inline()` reads per-line rects.
-            getClientRects: () => getReferenceClientRects(state.anchor),
+            getBoundingClientRect: () => {
+              const content = state.anchor.firstElementChild ?? state.anchor;
+              const rect = content.getBoundingClientRect();
+              const el =
+                rect.width || rect.height
+                  ? content
+                  : (content.firstElementChild ?? content);
+              return el.getBoundingClientRect();
+            },
+            getClientRects: () => {
+              const content = state.anchor.firstElementChild ?? state.anchor;
+              const rect = content.getBoundingClientRect();
+              const el =
+                rect.width || rect.height
+                  ? content
+                  : (content.firstElementChild ?? content);
+              return el.getClientRects();
+            },
           }
         : undefined,
     [state],
diff --git a/packages/react/src/editor/BlockNoteDefaultUI.tsx b/packages/react/src/editor/BlockNoteDefaultUI.tsx
index c85eef3973..75d618dc71 100644
--- a/packages/react/src/editor/BlockNoteDefaultUI.tsx
+++ b/packages/react/src/editor/BlockNoteDefaultUI.tsx
@@ -7,7 +7,6 @@ import {
   SuggestionMenu,
   TableHandlesExtension,
 } from "@blocknote/core/extensions";
-import { AttributionExtension } from "@blocknote/core/y";
 import { lazy, Suspense } from "react";
 
 import { FilePanelController } from "../components/FilePanel/FilePanelController.js";
@@ -162,7 +161,7 @@ export function BlockNoteDefaultUI(props: BlockNoteDefaultUIProps) {
           <FloatingThreadController portalElement={commentsPortal} />
         </Suspense>
       )}
-      {editor.getExtension(AttributionExtension) &&
+      {editor.getExtension("attribution") &&
         props.attributionTooltip !== false && (
           <AttributionTooltipController
             portalElement={attributionTooltipPortal}
diff --git a/packages/react/src/hooks/useExtension.ts b/packages/react/src/hooks/useExtension.ts
index d6e4fcd526..ec8c7a707f 100644
--- a/packages/react/src/hooks/useExtension.ts
+++ b/packages/react/src/hooks/useExtension.ts
@@ -49,13 +49,17 @@ export function useExtensionState<
     : never,
   TSelected = NoInfer<ExtractStore<TStore>>,
 >(
-  plugin: T,
+  plugin: T | string,
   ctx?: {
     editor?: BlockNoteEditor<any, any, any>;
     selector?: (state: NoInfer<ExtractStore<TStore>>) => TSelected;
   },
 ): TSelected {
-  const { store } = useExtension(plugin, ctx);
+  const extension = useExtension(
+    plugin as ExtensionFactory | Extension | string,
+    ctx,
+  );
+  const { store } = extension;
   if (!store) {
     throw new Error("Store not found on plugin", { cause: { plugin } });
   }