diff --git a/packages/text-editor/src/components/extension/nodeHighlight.ts b/packages/text-editor/src/components/extension/nodeHighlight.ts
index bbfa23e273..1b42b8fdb2 100644
--- a/packages/text-editor/src/components/extension/nodeHighlight.ts
+++ b/packages/text-editor/src/components/extension/nodeHighlight.ts
@@ -3,7 +3,7 @@ import { type Node as ProseMirrorNode, type MarkType } from '@tiptap/pm/model'
 import { Plugin, PluginKey, TextSelection } from '@tiptap/pm/state'
 import { AddMarkStep, RemoveMarkStep } from '@tiptap/pm/transform'
 import { Decoration, DecorationSet } from '@tiptap/pm/view'
-import { NodeUuidExtension, type NodeUuidOptions, type NodeUuidStorage, findNodeUuidMark } from './nodeUuid'
+import { NodeUuidExtension, type NodeUuidOptions, findNodeUuidMark } from './nodeUuid'
 import { type TextEditorCommand } from '../../types'
 
 export enum NodeHighlightType {
@@ -17,7 +17,7 @@ export function highlightUpdateCommand (): TextEditorCommand {
 }
 
 export interface NodeHighlightExtensionOptions extends NodeUuidOptions {
-  getNodeHighlightType: (uuid: string) => NodeHighlightType | undefined | null
+  getNodeHighlight: (uuid: string) => { type: NodeHighlightType, isActive?: boolean } | undefined | null
   isHighlightModeOn: () => boolean
   isAutoSelect?: () => boolean
 }
@@ -32,21 +32,23 @@ const generateAttributes = (uuid: string, options: NodeHighlightExtensionOptions
     return undefined
   }
 
-  const type = options.getNodeHighlightType(uuid)
-  if (type === null || type === undefined) {
+  const highlight = options.getNodeHighlight(uuid)
+  if (highlight === null || highlight === undefined) {
     return undefined
   }
   const classAttrs: { class?: string } = {}
 
-  if (type === NodeHighlightType.WARNING) {
+  if (highlight.type === NodeHighlightType.WARNING) {
     classAttrs.class = 'text-editor-highlighted-node-warning'
-  } else if (type === NodeHighlightType.ADD) {
+  } else if (highlight.type === NodeHighlightType.ADD) {
     classAttrs.class = 'text-editor-highlighted-node-add'
-  } else if (type === NodeHighlightType.DELETE) {
+  } else if (highlight.type === NodeHighlightType.DELETE) {
     classAttrs.class = 'text-editor-highlighted-node-delete'
   }
 
-  return classAttrs
+  return highlight.isActive === true
+    ? mergeAttributes(classAttrs, { class: 'text-editor-highlighted-node-selected' })
+    : classAttrs
 }
 
 const NodeHighlight = 'node-highlight'
@@ -69,17 +71,12 @@ declare module '@tiptap/core' {
 /**
  * Extension allows to highlight nodes based on uuid
  */
-export const NodeHighlightExtension: Extension<NodeHighlightExtensionOptions, NodeUuidStorage> =
+export const NodeHighlightExtension: Extension<NodeHighlightExtensionOptions> =
   Extension.create<NodeHighlightExtensionOptions>({
     name: NodeHighlight,
 
-    addStorage (): NodeUuidStorage {
-      return { activeNodeUuid: null }
-    },
-
     addProseMirrorPlugins () {
       const options = this.options
-      const storage: NodeUuidStorage = this.storage
 
       const plugins = [
         ...(this.parent?.() ?? []),
@@ -114,14 +111,14 @@ export const NodeHighlightExtension: Extension<NodeHighlightExtensionOptions, No
             init (_config, state): DecorationSet {
               const { doc, schema } = state
               const markType = schema.marks[NodeUuidExtension.name]
-              return createDecorations(doc, markType, storage, options)
+              return createDecorations(doc, markType, options)
             },
 
             apply (tr, decorations, oldState, newState) {
               const markType = newState.schema.marks[NodeUuidExtension.name]
 
               if (tr.getMeta(NodeHighlightMeta) !== undefined) {
-                return createDecorations(tr.doc, markType, storage, options)
+                return createDecorations(tr.doc, markType, options)
               }
 
               if (!tr.docChanged) {
@@ -136,7 +133,7 @@ export const NodeHighlightExtension: Extension<NodeHighlightExtensionOptions, No
                     (step instanceof RemoveMarkStep && step.mark.type === markType)
                 )
               ) {
-                return createDecorations(tr.doc, markType, storage, options)
+                return createDecorations(tr.doc, markType, options)
               }
 
               // update all decorations when changed content has mark changes
@@ -157,7 +154,7 @@ export const NodeHighlightExtension: Extension<NodeHighlightExtensionOptions, No
               })
 
               if (hasMarkChanges) {
-                return createDecorations(tr.doc, markType, storage, options)
+                return createDecorations(tr.doc, markType, options)
               }
 
               return decorations.map(tr.mapping, tr.doc)
@@ -189,17 +186,12 @@ export const NodeHighlightExtension: Extension<NodeHighlightExtensionOptions, No
 
     addExtensions () {
       const options: NodeHighlightExtensionOptions = this.options
-      const storage: NodeUuidStorage = this.storage
       return [
         NodeUuidExtension.extend({
           addOptions (): NodeUuidOptions {
             return {
               ...this.parent?.(),
-              ...options,
-              onNodeSelected: (uuid: string) => {
-                storage.activeNodeUuid = uuid
-                options.onNodeSelected?.(uuid)
-              }
+              ...options
             }
           }
         })
@@ -210,7 +202,6 @@ export const NodeHighlightExtension: Extension<NodeHighlightExtensionOptions, No
 const createDecorations = (
   doc: ProseMirrorNode,
   markType: MarkType,
-  storage: NodeUuidStorage,
   options: NodeHighlightExtensionOptions
 ): DecorationSet => {
   const decorations: Decoration[] = []
@@ -231,16 +222,7 @@ const createDecorations = (
         return
       }
 
-      decorations.push(
-        Decoration.inline(
-          range.from,
-          range.to,
-          mergeAttributes(
-            attributes,
-            nodeUuid === storage.activeNodeUuid ? { class: 'text-editor-highlighted-node-selected' } : {}
-          )
-        )
-      )
+      decorations.push(Decoration.inline(range.from, range.to, attributes))
     }
   })