diff --git a/wax-prosemirror-services/src/AiService/ToggleAiTool.js b/wax-prosemirror-services/src/AiService/ToggleAiTool.js
index 1cc0fa43feb72c280454b3b055dc36ab19ab69aa..88070ad692dee3befb64023e965e88ee893b6721 100644
--- a/wax-prosemirror-services/src/AiService/ToggleAiTool.js
+++ b/wax-prosemirror-services/src/AiService/ToggleAiTool.js
@@ -10,12 +10,7 @@ class ToggleAiTool extends Tools {
 
   renderTool(view) {
     return (
-      <ToggleAiComponent
-        displayed={false}
-        item={this.toJSON()}
-        key={uuidv4()}
-        view={view}
-      />
+      <ToggleAiComponent item={this.toJSON()} key={uuidv4()} view={view} />
     );
   }
 }
diff --git a/wax-prosemirror-services/src/AiService/components/AskAIOverlay.js b/wax-prosemirror-services/src/AiService/components/AskAIOverlay.js
index 5e5464b60b91b1e3dab00353e645acbadac3b5d7..4178becd06d09d74788c3976bff5a4e9301b0462 100644
--- a/wax-prosemirror-services/src/AiService/components/AskAIOverlay.js
+++ b/wax-prosemirror-services/src/AiService/components/AskAIOverlay.js
@@ -110,7 +110,7 @@ const AskAIOverlay = ({ setPosition, position, config }) => {
   const [result, setResult] = useState('');
   const [isSubmitted, setIsSubmitted] = useState(false);
   const [isLoading, setIsLoading] = useState(false);
-  const { AskAiContentTransformation, AiOn } = config;
+  const { AskAiContentTransformation } = config;
   const inputRef = useRef(null);
 
   useLayoutEffect(() => {
@@ -135,7 +135,7 @@ const AskAIOverlay = ({ setPosition, position, config }) => {
     }
 
     setPosition({ ...position, left, top });
-  }, [position.left, options.AiOn, AiOn]);
+  }, [position.left, options.AiOn]);
 
   const tryAgain = () => {
     // Reset the state to initial values
@@ -198,7 +198,7 @@ const AskAIOverlay = ({ setPosition, position, config }) => {
     }
   };
 
-  return options?.AiOn && AiOn ? (
+  return options?.AiOn ? (
     <Wrapper id="ai-overlay">
       <AskAIForm>
         <AskAIFormInput
diff --git a/wax-prosemirror-services/src/AiService/components/ToggleAiComponent.js b/wax-prosemirror-services/src/AiService/components/ToggleAiComponent.js
index c5954cba56b8f75bd418aecf055b30a279768618..066175e57b2da8ec5e2b9971735cd1b9491c0ddc 100644
--- a/wax-prosemirror-services/src/AiService/components/ToggleAiComponent.js
+++ b/wax-prosemirror-services/src/AiService/components/ToggleAiComponent.js
@@ -1,30 +1,39 @@
-import React, { useContext, useMemo } from 'react';
+import React, { useContext, useMemo, useState } from 'react';
 import { WaxContext, MenuButton } from 'wax-prosemirror-core';
 import PropTypes from 'prop-types';
 
 const ToggleAiComponent = ({ item }) => {
-  const { app, pmViews, options } = useContext(WaxContext);
-  const enableService = app.config.get('config.AskAiContentService');
+  const [checked, setChecked] = useState(false);
+  const context = useContext(WaxContext);
+  const {
+    pmViews: { main },
+  } = context;
 
   let isDisabled = false;
-  const { main } = pmViews;
   const isEditable = main.props.editable(editable => {
     return editable;
   });
 
   if (!isEditable) isDisabled = true;
 
+  const onMouseDown = () => {
+    context.setOption({ AiOn: !checked });
+    setChecked(!checked);
+    main.dispatch(main.state.tr.setMeta('addToHistory', false));
+    main.focus();
+  };
+
   return useMemo(
-    () =>
-      enableService.AiOn ? (
-        <MenuButton
-          active={false}
-          disabled={!isEditable || !options?.AiOn}
-          iconName={item.icon}
-          title={item.title}
-        />
-      ) : null,
-    [enableService.AiOn, isDisabled],
+    () => (
+      <MenuButton
+        active={checked}
+        disabled={!isEditable}
+        iconName={item.icon}
+        onMouseDown={onMouseDown}
+        title={item.title}
+      />
+    ),
+    [checked, isDisabled],
   );
 };