diff --git a/wax-prosemirror-services/src/AiService/plugins/AskAiSelectionPlugin.js b/wax-prosemirror-services/src/AiService/plugins/AskAiSelectionPlugin.js
index 41672bdf6104e2361b0adbfcb3be9a20101f3799..121981709ae853fdbba9bbb2cc1a7b1c088f3ae9 100644
--- a/wax-prosemirror-services/src/AiService/plugins/AskAiSelectionPlugin.js
+++ b/wax-prosemirror-services/src/AiService/plugins/AskAiSelectionPlugin.js
@@ -3,20 +3,28 @@ import { Decoration, DecorationSet } from 'prosemirror-view';
 
 const key = new PluginKey('askAiSelectionPlugin');
 
-export default props => {
+export default () => {
   return new Plugin({
     key,
     state: {
-      init: (_, state) => {
+      init: () => {
         return {};
       },
       apply(tr, prev, prevState, newState) {
         let createDecoration;
         const askAiInput = document.getElementById('askAiInput');
-
         if (askAiInput) {
+          const selectionWhenBlured = tr.getMeta(key);
+
+          const from = selectionWhenBlured
+            ? selectionWhenBlured.from
+            : newState.selection.from;
+          const to = selectionWhenBlured
+            ? selectionWhenBlured.to
+            : newState.selection.to;
+
           createDecoration = DecorationSet.create(newState.doc, [
-            Decoration.inline(newState.selection.from, newState.selection.to, {
+            Decoration.inline(from, to, {
               class: 'ask-ai-selection',
             }),
           ]);
@@ -32,14 +40,14 @@ export default props => {
         const askAiSelectionPluginState = state && key.getState(state);
         return askAiSelectionPluginState.createDecoration;
       },
-      //   handleDOMEvents: {
-      //     blur(view) {
-      //       view.dispatch(view.state.tr.setMeta(key, false));
-      //     },
-      //     focus(view) {
-      //       view.dispatch(view.state.tr.setMeta(key, true));
-      //     },
-      //   },
+      handleDOMEvents: {
+        blur(view) {
+          view.dispatch(view.state.tr.setMeta(key, view.state.selection));
+        },
+        //   focus(view) {
+        //     view.dispatch(view.state.tr.setMeta(key, true));
+        //   },
+      },
     },
   });
 };