diff --git a/package-lock.json b/package-lock.json
index 6fa1c8bb..607d7312 100644
--- a/package-lock.json
+++ b/package-lock.json
@@ -47,6 +47,7 @@
         "@tiptap/extension-bubble-menu": "^2.4.0",
         "@tiptap/extension-character-count": "^2.7.2",
         "@tiptap/extension-document": "^2.5.0",
+        "@tiptap/extension-highlight": "^2.9.1",
         "@tiptap/extension-link": "^2.2.4",
         "@tiptap/extension-paragraph": "^2.5.0",
         "@tiptap/extension-table": "^2.4.0",
@@ -5838,6 +5839,18 @@
         "@tiptap/core": "^2.0.0"
       }
     },
+    "node_modules/@tiptap/extension-highlight": {
+      "version": "2.9.1",
+      "resolved": "https://registry.npmjs.org/@tiptap/extension-highlight/-/extension-highlight-2.9.1.tgz",
+      "integrity": "sha512-ro3bARRgxb4v8w6fPVfG1kO2UWtLgKI5ESfsQ9CqiZuRkZdRKhM5ZpXPIky28Pn7CxhDUSXBXS/MhvP0VuhMJQ==",
+      "funding": {
+        "type": "github",
+        "url": "https://github.com/sponsors/ueberdosis"
+      },
+      "peerDependencies": {
+        "@tiptap/core": "^2.7.0"
+      }
+    },
     "node_modules/@tiptap/extension-history": {
       "version": "2.2.4",
       "license": "MIT",
diff --git a/package.json b/package.json
index aaf49a41..29efc44d 100644
--- a/package.json
+++ b/package.json
@@ -66,6 +66,7 @@
     "@tiptap/extension-bubble-menu": "^2.4.0",
     "@tiptap/extension-character-count": "^2.7.2",
     "@tiptap/extension-document": "^2.5.0",
+    "@tiptap/extension-highlight": "^2.9.1",
     "@tiptap/extension-link": "^2.2.4",
     "@tiptap/extension-paragraph": "^2.5.0",
     "@tiptap/extension-table": "^2.4.0",
diff --git a/src/components/Editor/EditorManager.tsx b/src/components/Editor/EditorManager.tsx
index e5b2b84c..18c0a999 100644
--- a/src/components/Editor/EditorManager.tsx
+++ b/src/components/Editor/EditorManager.tsx
@@ -1,7 +1,7 @@
 /* eslint-disable react/button-has-type */
 import React, { useEffect, useState } from 'react'
 import { EditorContent, BubbleMenu } from '@tiptap/react'
-import { getHTMLFromFragment } from '@tiptap/core'
+import { getHTMLFromFragment, Range } from '@tiptap/core'
 import TurndownService from 'turndown'
 import EditorContextMenu from './EditorContextMenu'
 import SearchBar from './Search/SearchBar'
@@ -17,6 +17,7 @@ const EditorManager: React.FC = () => {
   const [showAIPopup, setShowAIPopup] = useState(false)
   const { editor } = useFileContext()
   const turndownService = new TurndownService()
+  const [selectedRange, setSelectedRange] = useState<Range | null>(null)
 
   const handleContextMenu = (event: React.MouseEvent<HTMLDivElement>) => {
     event.preventDefault()
@@ -45,6 +46,23 @@ const EditorManager: React.FC = () => {
     window.ipcRenderer.on('editor-flex-center-changed', handleEditorChange)
   }, [])
 
+  useEffect(() => {
+    if (!editor) return
+
+    if (showAIPopup && selectedRange) {
+      editor.chain().focus().setMark('highlight').run()
+    }
+  }, [showAIPopup, selectedRange, editor])
+
+  useEffect(() => {
+    if (!editor) return
+
+    if (!showAIPopup) {
+      editor.chain().focus().unsetMark('highlight').run()
+      editor.chain().focus().toggleHighlight().run()
+    }
+  }, [showAIPopup, editor])
+
   return (
     <div
       className="relative size-full cursor-text overflow-hidden bg-dark-gray-c-eleven py-4 text-slate-400 opacity-80"
@@ -61,6 +79,8 @@ const EditorManager: React.FC = () => {
             interactiveBorder: 20,
             onHidden: () => {
               setShowAIPopup(false)
+              console.log('unsetting selected range')
+              setSelectedRange(null)
             },
             maxWidth: 'none',
           }}
@@ -71,6 +91,17 @@ const EditorManager: React.FC = () => {
               e.preventDefault()
               e.stopPropagation()
             }}
+            onMouseUp={(e) => {
+              e.preventDefault()
+              e.stopPropagation()
+
+              if (!showAIPopup) {
+                setSelectedRange({
+                  from: editor.state.selection.from,
+                  to: editor.state.selection.to,
+                })
+              }
+            }}
             onClick={(e) => {
               e.preventDefault()
               e.stopPropagation()
@@ -80,7 +111,10 @@ const EditorManager: React.FC = () => {
               <AiEditMenu
                 selectedText={turndownService.turndown(
                   getHTMLFromFragment(
-                    editor.state.doc.slice(editor.state.selection.from, editor.state.selection.to).content,
+                    editor.state.doc.slice(
+                      selectedRange?.from || editor.state.selection.from,
+                      selectedRange?.to || editor.state.selection.to,
+                    ).content,
                     editor.schema,
                   ),
                 )}
@@ -89,11 +123,12 @@ const EditorManager: React.FC = () => {
                     .chain()
                     .focus()
                     .deleteRange({
-                      from: editor.state.selection.from,
-                      to: editor.state.selection.to,
+                      from: selectedRange?.from || editor.state.selection.from,
+                      to: selectedRange?.to || editor.state.selection.to,
                     })
                     .insertContent(newText)
                     .run()
+                  setShowAIPopup(false)
                 }}
               />
             ) : (
diff --git a/src/contexts/FileContext.tsx b/src/contexts/FileContext.tsx
index 87a02965..db8a25d2 100644
--- a/src/contexts/FileContext.tsx
+++ b/src/contexts/FileContext.tsx
@@ -19,6 +19,7 @@ import { toast } from 'react-toastify'
 import { Markdown } from 'tiptap-markdown'
 import { useDebounce } from 'use-debounce'
 import { FileInfo, FileInfoTree } from 'electron/main/filesystem/types'
+import Highlight from '@tiptap/extension-highlight'
 import {
   findRelevantDirectoriesToBeExpanded,
   flattenFileInfoTree,
@@ -181,6 +182,7 @@ export const FileProvider: React.FC<{ children: ReactNode }> = ({ children }) =>
         transformPastedText: true,
         transformCopiedText: false,
       }),
+      Highlight,
       TaskItem.configure({
         nested: true,
       }),
diff --git a/src/styles/global.css b/src/styles/global.css
index aeafb8aa..a32f9334 100644
--- a/src/styles/global.css
+++ b/src/styles/global.css
@@ -190,6 +190,14 @@ button {
   /* Adjust day cell size */
 }
 
+/* AI Edit highlight styles */
+/* .tiptap .ai-edit-highlight {
+  background-color: pink; 
+  border-radius: 2px;
+  padding: 2px 0;
+  margin: 0 -2px;
+} */
+
 @layer base {
   :root {
     --background: 0 0% 100%;