diff --git a/assets/javascripts/discourse/connectors/post-text-buttons/ai-helper-options-menu.gjs b/assets/javascripts/discourse/connectors/post-text-buttons/ai-helper-options-menu.gjs
index f76820e7..fce1b497 100644
--- a/assets/javascripts/discourse/connectors/post-text-buttons/ai-helper-options-menu.gjs
+++ b/assets/javascripts/discourse/connectors/post-text-buttons/ai-helper-options-menu.gjs
@@ -100,8 +100,13 @@ export default class AIHelperOptionsMenu extends Component {
     );
 
     const textNodes = [];
-    while (walker.nextNode()) {
+
+    if (walker.currentNode?.nodeType === Node.TEXT_NODE) {
       textNodes.push(walker.currentNode);
+    } else {
+      while (walker.nextNode()) {
+        textNodes.push(walker.currentNode);
+      }
     }
 
     for (let textNode of textNodes) {
diff --git a/assets/stylesheets/modules/ai-helper/common/ai-helper.scss b/assets/stylesheets/modules/ai-helper/common/ai-helper.scss
index 7ad2d43c..0c97b227 100644
--- a/assets/stylesheets/modules/ai-helper/common/ai-helper.scss
+++ b/assets/stylesheets/modules/ai-helper/common/ai-helper.scss
@@ -142,7 +142,7 @@
 }
 
 .ai-helper-highlighted-selection {
-  background-color: var(--highlight-low);
+  background-color: var(--highlight-low-or-medium);
 }
 
 // AI Typing indicator (taken from: https://github.com/nzbin/three-dots)