1 files changed, 39 insertions, 4 deletions
diff --git a/app/assets/javascripts/content_editor/extensions/code_block_highlight.js b/app/assets/javascripts/content_editor/extensions/code_block_highlight.js
index 204ac07d401..61f379fc0a2 100644
--- a/app/assets/javascripts/content_editor/extensions/code_block_highlight.js
+++ b/app/assets/javascripts/content_editor/extensions/code_block_highlight.js
@@ -1,10 +1,21 @@
 import { CodeBlockLowlight } from '@tiptap/extension-code-block-lowlight';
-import { lowlight } from 'lowlight/lib/all';
+import { textblockTypeInputRule } from '@tiptap/core';
+import codeBlockLanguageLoader from '../services/code_block_language_loader';
 
 const extractLanguage = (element) => element.getAttribute('lang');
+export const backtickInputRegex = /^```([a-z]+)?[\s\n]$/;
+export const tildeInputRegex = /^~~~([a-z]+)?[\s\n]$/;
 
 export default CodeBlockLowlight.extend({
   isolating: true,
+  exitOnArrowDown: false,
+
+  addOptions() {
+    return {
+      ...this.parent?.(),
+      languageLoader: codeBlockLanguageLoader,
+    };
+  },
 
   addAttributes() {
     return {
@@ -18,16 +29,40 @@ export default CodeBlockLowlight.extend({
       },
     };
   },
+  addInputRules() {
+    const { languageLoader } = this.options;
+    const getAttributes = (match) => languageLoader?.loadLanguageFromInputRule(match) || {};
+
+    return [
+      textblockTypeInputRule({
+        find: backtickInputRegex,
+        type: this.type,
+        getAttributes,
+      }),
+      textblockTypeInputRule({
+        find: tildeInputRegex,
+        type: this.type,
+        getAttributes,
+      }),
+    ];
+  },
+  parseHTML() {
+    return [
+      ...(this.parent?.() || []),
+      {
+        tag: 'div.markdown-code-block',
+        skip: true,
+      },
+    ];
+  },
   renderHTML({ HTMLAttributes }) {
     return [
       'pre',
       {
         ...HTMLAttributes,
-        class: `content-editor-code-block ${HTMLAttributes.class}`,
+        class: `content-editor-code-block ${gon.user_color_scheme} ${HTMLAttributes.class}`,
       },
       ['code', {}, 0],
     ];
   },
-}).configure({
-  lowlight,
 });