diff --git a/packages/tiptap-extensions/src/nodes/CodeBlockHighlight.js b/packages/tiptap-extensions/src/nodes/CodeBlockHighlight.js
index 76c690a0..99c1b417 100644
--- a/packages/tiptap-extensions/src/nodes/CodeBlockHighlight.js
+++ b/packages/tiptap-extensions/src/nodes/CodeBlockHighlight.js
@@ -122,13 +122,16 @@ export default class CodeBlockHighlight extends Node {
           init(_, { doc }) {
             return getDecorations(doc)
           },
-          apply(tr, set) {
+          apply(transaction, decorationSet) {
             // TODO: find way to cache decorations
             // see: https://discuss.prosemirror.net/t/how-to-update-multiple-inline-decorations-on-node-change/1493
-            if (tr.docChanged) {
-              return getDecorations(tr.doc)
+
+            const currentNodeName = transaction.curSelection.$head.parent.type.name
+            if (transaction.docChanged && currentNodeName === 'code_block') {
+              return getDecorations(transaction.doc)
             }
-            return set.map(tr.mapping, tr.doc)
+
+            return decorationSet.map(transaction.mapping, transaction.doc)
           },
         },
         props: {