diff --git a/packages/core/src/plugins/codeBlockWrapButtons.ts b/packages/core/src/plugins/codeBlockWrapButtons.ts
index 4d98c61319..3ebe226cc3 100644
--- a/packages/core/src/plugins/codeBlockWrapButtons.ts
+++ b/packages/core/src/plugins/codeBlockWrapButtons.ts
@@ -37,16 +37,28 @@ function getButtonHTML() {
 }
 
 const wrapCodeBlockScript = `<script>
-    function toggleCodeBlockWrap(element) {
-      const pre = element.parentElement.parentElement;
-      const classList = pre.querySelector('code').classList; 
-      if (classList.contains('wrap')) {
-          classList.remove('wrap');
-      } else {
-          classList.add('wrap')
-      }
+function toggleCodeBlockWrap(element) {
+  const pre = element.parentElement.parentElement;
+  const code = pre.querySelector('code');
+  const classList = code.classList; 
+
+  if (code.dataset.originalLineNumbers === undefined) {
+    code.dataset.originalLineNumbers = classList.contains('line-numbers').toString();
+  }
+
+  if (classList.contains('wrap')) {
+    classList.remove('wrap');
+    if (code.dataset.originalLineNumbers === 'false') {
+      classList.remove('line-numbers');
     }
-    </script>`;
+  } else {
+    classList.add('wrap');
+    if (code.dataset.originalLineNumbers === 'false') {
+      classList.add('line-numbers');
+    }
+  }
+}
+</script>`;
 
 export = {
   getScripts: () => [wrapCodeBlockScript],