Fix code block syntax highlighting not updating when changing its language

packages/editor/src/extensions/code-block/code-block.ts

@@ -127,7 +127,7 @@ export const CodeBlock = Node.create<CodeBlockOptions>({
       id: {
         default: undefined,
         rendered: false,
-        parseHTML: () => `codeblock-${nanoid(12)}`
+        parseHTML: () => createCodeblockId()
       },
       caretPosition: {
         default: undefined,
@@ -227,13 +227,18 @@ export const CodeBlock = Node.create<CodeBlockOptions>({
       setCodeBlock:
         (attributes) =>
         ({ commands }) => {
-          return commands.setNode(this.name, attributes);
+          return commands.setNode(this.name, {
+            ...attributes,
+            id: createCodeblockId()
+          });
         },
       toggleCodeBlock:
         (attributes) =>
         ({ commands }) => {
-          console.log("TOGGLING!");
-          return commands.toggleNode(this.name, "paragraph", attributes);
+          return commands.toggleNode(this.name, "paragraph", {
+            ...attributes,
+            id: createCodeblockId()
+          });
         },
       changeCodeBlockIndentation:
         (options) =>
@@ -477,14 +482,16 @@ export const CodeBlock = Node.create<CodeBlockOptions>({
         find: backtickInputRegex,
         type: this.type,
         getAttributes: (match) => ({
-          language: match[1]
+          language: match[1],
+          id: createCodeblockId()
         })
       }),
       textblockTypeInputRule({
         find: tildeInputRegex,
         type: this.type,
         getAttributes: (match) => ({
-          language: match[1]
+          language: match[1],
+          id: createCodeblockId()
         })
       })
     ];
@@ -520,7 +527,7 @@ export const CodeBlock = Node.create<CodeBlockOptions>({
             if (isCode && !isInsideCodeBlock) {
               tr.replaceSelectionWith(
                 this.type.create({
-                  id: `codeblock-${nanoid(12)}`,
+                  id: createCodeblockId(),
                   language,
                   indentType: indent.type,
                   indentLength: indent.amount
@@ -805,3 +812,7 @@ export function inferLanguage(node: Element) {
   );
   return language?.filename;
 }
+
+function createCodeblockId() {
+  return `codeblock-${nanoid(12)}`;
+}

packages/editor/src/extensions/code-block/highlighter.ts

@@ -125,9 +125,14 @@ export function HighlighterPlugin({
 
           const changedBlocks: Set<string> = new Set();
           for (const blockKey in pluginState.languages) {
-            if (HIGHLIGHTED_BLOCKS.has(blockKey)) continue;
-
             const language = pluginState.languages[blockKey];
+            if (
+              HIGHLIGHTED_BLOCKS.has(blockKey) &&
+              refractor.registered(language)
+            ) {
+              continue;
+            }
+
             const languageDefinition = Languages.find(
               (l) =>
                 l.filename === language || l.alias?.some((a) => a === language)
@@ -193,12 +198,21 @@ export function HighlighterPlugin({
           });
           if (changedBlocks.length > 0) {
             const updated: Set<number> = new Set();
+            let hasChanges = false;
+
             changedBlocks.forEach((block) => {
               if (updated.has(block.pos)) return;
               updated.add(block.pos);
 
               const { id, language } = block.node.attrs;
-              if (languages[id]) {
+
+              if (
+                !languages[id] ||
+                (language && !refractor.registered(language))
+              ) {
+                languages[id] = language;
+                hasChanges = true;
+              } else {
                 const newDecorations = getDecorations({
                   block,
                   defaultLanguage
@@ -221,12 +235,12 @@ export function HighlighterPlugin({
                   decorations = decorations.remove(toRemove);
                 if (toAdd.length > 0)
                   decorations = decorations.add(tr.doc, toAdd);
-              } else {
-                languages[id] = language;
+
+                hasChanges = true;
               }
             });
 
-            if (decorations !== pluginState.decorations) {
+            if (hasChanges) {
               return { decorations, languages };
             }
           }

packages/editor/src/extensions/code-block/tests/__snapshots__/code-block.test.ts.snap

@@ -1,5 +1,7 @@
 // Vitest Snapshot v1, https://vitest.dev/guide/snapshot.html
 
+exports[`Adding a new codeblock & changing the language should apply the new highlighting 1`] = `"<div><div contenteditable=\\"true\\" translate=\\"no\\" tabindex=\\"0\\" class=\\"ProseMirror\\"><div class=\\"codeblock-view-content-wrap\\"><div class=\\"node-content-wrapper\\" style=\\"white-space: pre; min-width: 20px;\\"><span class=\\"token keyword\\">function</span> <span class=\\"token function\\">hello</span><span class=\\"token punctuation\\">(</span><span class=\\"token punctuation\\">)</span> <span class=\\"token punctuation\\">{</span> <span class=\\"token punctuation\\">}</span></div></div></div></div>"`;
+
 exports[`codeblocks should get highlighted after pasting 1`] = `"<div><div contenteditable=\\"true\\" translate=\\"no\\" tabindex=\\"0\\" class=\\"ProseMirror\\"><div class=\\"codeblock-view-content-wrap\\"><div class=\\"node-content-wrapper\\" style=\\"white-space: pre; min-width: 20px;\\"><span class=\\"token keyword\\">function</span> <span class=\\"token function\\">hello</span><span class=\\"token punctuation\\">(</span><span class=\\"token punctuation\\">)</span> <span class=\\"token punctuation\\">{</span> <span class=\\"token punctuation\\">}</span></div></div><div class=\\"codeblock-view-content-wrap\\"><div class=\\"node-content-wrapper\\" style=\\"white-space: pre; min-width: 20px;\\"><span class=\\"token keyword\\">function</span> <span class=\\"token function\\">hello</span><span class=\\"token punctuation\\">(</span><span class=\\"token punctuation\\">)</span> <span class=\\"token punctuation\\">{</span> <span class=\\"token punctuation\\">}</span></div></div></div></div>"`;
 
 exports[`codeblocks should get highlighted on init 1`] = `"<div><div contenteditable=\\"true\\" translate=\\"no\\" tabindex=\\"0\\" class=\\"ProseMirror\\"><div class=\\"codeblock-view-content-wrap\\"><div class=\\"node-content-wrapper\\" style=\\"white-space: pre; min-width: 20px;\\"><span class=\\"token keyword\\">function</span> <span class=\\"token function\\">hello</span><span class=\\"token punctuation\\">(</span><span class=\\"token punctuation\\">)</span> <span class=\\"token punctuation\\">{</span> <span class=\\"token punctuation\\">}</span></div></div><div class=\\"codeblock-view-content-wrap\\"><div class=\\"node-content-wrapper\\" style=\\"white-space: pre; min-width: 20px;\\"><span class=\\"token keyword\\">function</span> <span class=\\"token function\\">hello</span><span class=\\"token punctuation\\">(</span><span class=\\"token punctuation\\">)</span> <span class=\\"token punctuation\\">{</span> <span class=\\"token punctuation\\">}</span></div></div></div></div>"`;

packages/editor/src/extensions/code-block/tests/code-block.test.ts

@@ -22,6 +22,7 @@ import { expect, test, vi } from "vitest";
 import { CodeBlock, inferLanguage } from "../code-block";
 import { HighlighterPlugin } from "../highlighter";
 import { getChangedNodes } from "../../../utils/prosemirror";
+import { refractor } from "refractor/lib/core";
 
 const CODEBLOCKS_HTML = h("div", [
   h("pre", [h("code", ["function hello() { }"])], {
@@ -191,3 +192,37 @@ test("editing code in a highlighted code block should not be too slow", async ()
   expect(timings.reduce((a, b) => a + b) / timings.length).toBeLessThan(16);
   expect(editorElement.outerHTML).toMatchSnapshot();
 });
+
+test("Adding a new codeblock & changing the language should apply the new highlighting", async () => {
+  const editorElement = h("div");
+  const { editor } = createEditor({
+    element: editorElement,
+    extensions: {
+      codeblock: CodeBlock
+    }
+  });
+
+  editor.commands.setCodeBlock();
+  editor.commands.insertContent("function hello() { }");
+
+  editor.commands.updateAttributes(CodeBlock.name, { language: "javascript" });
+
+  await new Promise((resolve) => setTimeout(resolve, 100));
+
+  expect(editorElement.outerHTML).toMatchSnapshot();
+  expect(refractor.registered("javascript")).toBe(true);
+});
+
+test("Switching codeblock language should register the new language", async () => {
+  const editorElement = h("div");
+  const { editor } = createEditor({
+    element: editorElement,
+    initialContent: CODEBLOCKS_HTML,
+    extensions: {
+      codeblock: CodeBlock
+    }
+  });
+  editor.commands.updateAttributes(CodeBlock.name, { language: "java" });
+  await new Promise((resolve) => setTimeout(resolve, 100));
+  expect(refractor.registered("java")).toBe(true);
+});