fix(ai): make sure ShowSelection works

Author: YousefED

packages/core/src/comments/extension.ts

@@ -306,15 +306,19 @@ export const CommentsExtension = createExtension(
           selectedThreadId: undefined,
           pendingComment: true,
         }));
-        editor.getExtension(ShowSelectionExtension)?.showSelection(true);
+        editor
+          .getExtension(ShowSelectionExtension)
+          ?.showSelection(true, "comments");
       },
       stopPendingComment() {
         store.setState((prev) => ({
           ...prev,
           selectedThreadId: undefined,
           pendingComment: false,
         }));
-        editor.getExtension(ShowSelectionExtension)?.showSelection(false);
+        editor
+          .getExtension(ShowSelectionExtension)
+          ?.showSelection(false, "comments");
       },
       async createThread(options: {
         initialComment: { body: CommentBody; metadata?: any };

packages/core/src/extensions/ShowSelection/ShowSelection.ts

@@ -14,7 +14,7 @@ const PLUGIN_KEY = new PluginKey(`blocknote-show-selection`);
  */
 export const ShowSelectionExtension = createExtension(({ editor }) => {
   const store = createStore(
-    { enabled: false },
+    { enabledSet: new Set<string>() },
     {
       onUpdate() {
         editor.transact((tr) => tr.setMeta(PLUGIN_KEY, {}));
@@ -30,7 +30,7 @@ export const ShowSelectionExtension = createExtension(({ editor }) => {
         props: {
           decorations: (state) => {
             const { doc, selection } = state;
-            if (!store.state.enabled) {
+            if (store.state.enabledSet.size === 0) {
               return DecorationSet.empty;
             }
             const dec = Decoration.inline(selection.from, selection.to, {
@@ -43,9 +43,19 @@ export const ShowSelectionExtension = createExtension(({ editor }) => {
     ],
     /**
      * Show or hide the selection decoration
+     *
+     * @param shouldShow - Whether to show the selection decoration
+     * @param key - The key of the selection to show or hide,
+     * this is necessary to prevent disabling ShowSelection from one place
+     * will interfere with other parts of the code that need to show the selection decoration
+     * (e.g.: CreateLinkButton and AIExtension)
      */
-    showSelection(shouldShow: boolean) {
-      store.setState({ enabled: shouldShow });
+    showSelection(shouldShow: boolean, key: string) {
+      store.setState({
+        enabledSet: shouldShow
+          ? new Set([...store.state.enabledSet, key])
+          : new Set([...store.state.enabledSet].filter((k) => k !== key)),
+      });
     },
   } as const;
 });

packages/react/src/components/FormattingToolbar/DefaultButtons/CreateLinkButton.tsx

@@ -49,8 +49,8 @@ export const CreateLinkButton = () => {
 
   const [showPopover, setShowPopover] = useState(false);
   useEffect(() => {
-    showSelection(showPopover);
-    return () => showSelection(false);
+    showSelection(showPopover, "createLinkButton");
+    return () => showSelection(false, "createLinkButton");
   }, [showPopover, showSelection]);
 
   const state = useEditorState({

packages/xl-ai/src/AIExtension.ts

@@ -143,7 +143,9 @@ export const AIExtension = createExtension(
        * Open the AI menu at a specific block
        */
       openAIMenuAtBlock(blockID: string) {
-        editor.getExtension(ShowSelectionExtension)?.showSelection(true);
+        editor
+          .getExtension(ShowSelectionExtension)
+          ?.showSelection(true, "aiMenu");
         editor.isEditable = false;
         store.setState({
           aiMenuState: {
@@ -167,7 +169,9 @@ export const AIExtension = createExtension(
           aiMenuState: "closed",
         });
         chatSession = undefined;
-        editor.getExtension(ShowSelectionExtension)?.showSelection(false);
+        editor
+          .getExtension(ShowSelectionExtension)
+          ?.showSelection(false, "aiMenu");
         editor.isEditable = true;
         editor.focus();
       },
@@ -337,7 +341,9 @@ export const AIExtension = createExtension(
         }
 
         if (status === "ai-writing") {
-          editor.getExtension(ShowSelectionExtension)?.showSelection(false);
+          editor
+            .getExtension(ShowSelectionExtension)
+            ?.showSelection(false, "aiMenu");
         }
 
         if (typeof status === "object") {