diff --git a/src/BubbleMenu.tsx b/src/BubbleMenu.tsx
index 3f48a2da966..415bb7023ff 100644
--- a/src/BubbleMenu.tsx
+++ b/src/BubbleMenu.tsx
@@ -1,5 +1,5 @@
 import React, { useEffect, useRef } from 'react'
-import { BubbleMenuPlugin, BubbleMenuPluginKey, BubbleMenuPluginProps } from '@tiptap/extension-bubble-menu'
+import { BubbleMenuPlugin, BubbleMenuPluginProps } from '@tiptap/extension-bubble-menu'
 
 export type BubbleMenuProps = Omit<BubbleMenuPluginProps, 'element'> & {
   className?: string,
@@ -9,16 +9,23 @@ export const BubbleMenu: React.FC<BubbleMenuProps> = props => {
   const element = useRef<HTMLDivElement>(null)
 
   useEffect(() => {
-    const { editor, tippyOptions } = props
+    const {
+      key,
+      editor,
+      tippyOptions,
+      shouldShow,
+    } = props
 
     editor.registerPlugin(BubbleMenuPlugin({
+      key,
       editor,
       element: element.current as HTMLElement,
       tippyOptions,
+      shouldShow,
     }))
 
     return () => {
-      editor.unregisterPlugin(BubbleMenuPluginKey)
+      editor.unregisterPlugin(key)
     }
   }, [])
 
diff --git a/src/FloatingMenu.tsx b/src/FloatingMenu.tsx
index e33a06363b8..d1b5d34c5fb 100644
--- a/src/FloatingMenu.tsx
+++ b/src/FloatingMenu.tsx
@@ -1,5 +1,5 @@
 import React, { useEffect, useRef } from 'react'
-import { FloatingMenuPlugin, FloatingMenuPluginKey, FloatingMenuPluginProps } from '@tiptap/extension-floating-menu'
+import { FloatingMenuPlugin, FloatingMenuPluginProps } from '@tiptap/extension-floating-menu'
 
 export type FloatingMenuProps = Omit<FloatingMenuPluginProps, 'element'> & {
   className?: string,
@@ -9,16 +9,23 @@ export const FloatingMenu: React.FC<FloatingMenuProps> = props => {
   const element = useRef<HTMLDivElement>(null)
 
   useEffect(() => {
-    const { editor, tippyOptions } = props
+    const {
+      key,
+      editor,
+      tippyOptions,
+      shouldShow,
+    } = props
 
     editor.registerPlugin(FloatingMenuPlugin({
+      key,
       editor,
       element: element.current as HTMLElement,
       tippyOptions,
+      shouldShow,
     }))
 
     return () => {
-      editor.unregisterPlugin(FloatingMenuPluginKey)
+      editor.unregisterPlugin(key)
     }
   }, [])