diff --git a/wax-prosemirror-services/src/AiService/AskAiContentService.js b/wax-prosemirror-services/src/AiService/AskAiContentService.js
index 6f79d22a051222af24397aeb5bf24be297995f73..466e9fc8cfb6312fe1ce5a395e5802bdf8266563 100644
--- a/wax-prosemirror-services/src/AiService/AskAiContentService.js
+++ b/wax-prosemirror-services/src/AiService/AskAiContentService.js
@@ -15,19 +15,15 @@ class AskAiContentService extends Service {
     );
 
     const createOverlay = this.container.get('CreateOverlay');
-    const { config } = this;
+    const config = this.config;
 
     // Create the overlay
-    createOverlay(
-      AskAIOverlay,
-      { config },
-      {
-        nodeType: '',
-        markType: '',
-        followCursor: false,
-        selection: true,
-      },
-    );
+    createOverlay(AskAIOverlay, config, {
+      nodeType: '',
+      markType: '',
+      followCursor: false,
+      selection: true,
+    });
   }
 
   register() {
diff --git a/wax-prosemirror-services/src/AiService/ReplaceSelectedText.js b/wax-prosemirror-services/src/AiService/ReplaceSelectedText.js
index e4a55dba4ec927e4761511248f8ba546cea4e25c..0fed7eec90f64c1313fa960275e70b529bdbedad 100644
--- a/wax-prosemirror-services/src/AiService/ReplaceSelectedText.js
+++ b/wax-prosemirror-services/src/AiService/ReplaceSelectedText.js
@@ -8,6 +8,7 @@ const elementFromString = string => {
 };
 
 const replaceSelectedText = (view, responseText, replace = false) => {
+  if (!view) return;
   let { state } = view;
   let { tr } = state;
   const { from, to } = tr.selection;
diff --git a/wax-prosemirror-services/src/AiService/components/AskAIOverlay.js b/wax-prosemirror-services/src/AiService/components/AskAIOverlay.js
index bc5227e4c3fd50b1b0189503c9810492f25b80d9..a7a7927ddeac2a78f02b59f5cc4fb1c2d38c1d7a 100644
--- a/wax-prosemirror-services/src/AiService/components/AskAIOverlay.js
+++ b/wax-prosemirror-services/src/AiService/components/AskAIOverlay.js
@@ -192,7 +192,7 @@ const AskAIOverlay = ({ setPosition, position, config }) => {
     }
   };
 
-  return options.AiOn ? (
+  return options?.AiOn ? (
     <Wrapper id="ai-overlay">
       <AskAIForm>
         <AskAIFormInput
diff --git a/wax-prosemirror-services/src/AiService/plugins/AskAiSelectionPlugin.js b/wax-prosemirror-services/src/AiService/plugins/AskAiSelectionPlugin.js
index 56ffbe9f184a1b7705df430cd9319e1d4355f0b5..ed7a8e29c64394560b040221c5ca9349c31621a3 100644
--- a/wax-prosemirror-services/src/AiService/plugins/AskAiSelectionPlugin.js
+++ b/wax-prosemirror-services/src/AiService/plugins/AskAiSelectionPlugin.js
@@ -14,10 +14,14 @@ export default () => {
         let createDecoration;
         const askAiInput = document.getElementById('askAiInput');
         if (askAiInput) {
-          // const selectionWhenBlured = tr.getMeta(key);
+          const selectionWhenBlured = tr.getMeta(key);
 
-          const from = newState.selection.from;
-          const to = newState.selection.to;
+          const from = selectionWhenBlured
+            ? selectionWhenBlured.from
+            : newState.selection.from;
+          const to = selectionWhenBlured
+            ? selectionWhenBlured.to
+            : newState.selection.to;
 
           createDecoration = DecorationSet.create(newState.doc, [
             Decoration.inline(from, to, {
@@ -38,7 +42,8 @@ export default () => {
       },
       handleDOMEvents: {
         blur(view) {
-          view.dispatch(view.state.tr.setMeta(key, view.state.selection));
+          if (view)
+            view.dispatch(view.state.tr.setMeta(key, view.state.selection));
         },
         //   focus(view) {
         //     view.dispatch(view.state.tr.setMeta(key, true));