fix codeblock/mermaid gap cursor

packages/editor-ext/src/lib/custom-code-block/custom-code-block.ts

@@ -1,5 +1,7 @@
 import type { CodeBlockOptions } from '@tiptap/extension-code-block';
 import CodeBlock from '@tiptap/extension-code-block';
+import { Plugin, Selection, TextSelection } from '@tiptap/pm/state';
+import { GapCursor } from '@tiptap/pm/gapcursor';
 
 import { LowlightPlugin } from './lowlight-plugin.js';
 import { ReactNodeViewRenderer } from '@tiptap/react';
@@ -19,7 +21,11 @@ const TAB_CHAR = '\u00A0\u00A0';
  * @see https://tiptap.dev/api/nodes/code-block-lowlight
  */
 export const CustomCodeBlock = CodeBlock.extend<CodeBlockLowlightOptions>({
+  // Run ahead of Gapcursor (100) so the mermaid arrow-into-source plugin
+  // can intercept before gapcursor takes over.
+  priority: 101,
   selectable: true,
+  isolating: true,
 
   addOptions() {
     return {
@@ -35,8 +41,86 @@ export const CustomCodeBlock = CodeBlock.extend<CodeBlockLowlightOptions>({
   },
 
   addKeyboardShortcuts() {
+    const isMermaid = (node: any) =>
+      node?.type === this.type && node.attrs.language === 'mermaid';
+
     return {
       ...this.parent?.(),
+      // Stop at the gap (or enter mermaid source) instead of jumping
+      // straight into the next block, so the user can place a cursor
+      // between two adjacent isolating blocks.
+      ArrowDown: ({ editor }) => {
+        const { state } = editor;
+        const { selection, doc } = state;
+        const { $from, empty } = selection;
+
+        if (!empty || $from.parent.type !== this.type) return false;
+        if ($from.parentOffset !== $from.parent.nodeSize - 2) return false;
+
+        const after = $from.after();
+        if (after >= doc.content.size) {
+          return editor.commands.exitCode();
+        }
+
+        const $after = doc.resolve(after);
+        const nodeAfter = $after.nodeAfter;
+
+        if (isMermaid(nodeAfter)) {
+          return editor.commands.command(({ tr }) => {
+            tr.setSelection(TextSelection.create(tr.doc, after + 1));
+            return true;
+          });
+        }
+
+        if (
+          nodeAfter?.type.spec.isolating &&
+          !nodeAfter.type.spec.atom
+        ) {
+          return editor.commands.command(({ tr }) => {
+            tr.setSelection(new GapCursor(tr.doc.resolve(after)));
+            return true;
+          });
+        }
+
+        return editor.commands.command(({ tr }) => {
+          tr.setSelection(Selection.near(tr.doc.resolve(after)));
+          return true;
+        });
+      },
+      // Mirror of ArrowDown; upstream has no ArrowUp handler.
+      ArrowUp: ({ editor }) => {
+        const { state } = editor;
+        const { selection, doc } = state;
+        const { $from, empty } = selection;
+
+        if (!empty || $from.parent.type !== this.type) return false;
+        if ($from.parentOffset !== 0) return false;
+
+        const before = $from.before();
+        if (before <= 0) return false;
+
+        const $before = doc.resolve(before);
+        const nodeBefore = $before.nodeBefore;
+
+        if (isMermaid(nodeBefore)) {
+          return editor.commands.command(({ tr }) => {
+            tr.setSelection(TextSelection.create(tr.doc, before - 1));
+            return true;
+          });
+        }
+
+        if (
+          nodeBefore?.type.spec.isolating &&
+          !nodeBefore.type.spec.atom
+        ) {
+          return editor.commands.command(({ tr }) => {
+            tr.setSelection(new GapCursor(tr.doc.resolve(before)));
+            return true;
+          });
+        }
+
+        return false;
+      },
       'Mod-a': () => {
         if (this.editor.isActive('codeBlock')) {
           const { state } = this.editor;
@@ -84,6 +168,7 @@ export const CustomCodeBlock = CodeBlock.extend<CodeBlockLowlightOptions>({
   },
 
   addProseMirrorPlugins() {
+    const codeBlockType = this.type;
     return [
       ...(this.parent?.() || []),
       LowlightPlugin({
@@ -91,6 +176,60 @@ export const CustomCodeBlock = CodeBlock.extend<CodeBlockLowlightOptions>({
         lowlight: this.options.lowlight,
         defaultLanguage: this.options.defaultLanguage,
       }),
+      // Mermaid hides its <pre> when unselected, so the browser's native
+      // vertical caret movement skips past it. Land the cursor inside the
+      // source explicitly.
+      new Plugin({
+        props: {
+          handleKeyDown: (view, event) => {
+            if (event.key !== 'ArrowUp' && event.key !== 'ArrowDown') {
+              return false;
+            }
+            const { state } = view;
+            const { selection } = state;
+            if (
+              !selection.empty ||
+              !(selection instanceof TextSelection)
+            ) {
+              return false;
+            }
+            const { $from } = selection;
+            if ($from.depth === 0 || $from.parent.type === codeBlockType) {
+              return false;
+            }
+            const dir = event.key === 'ArrowUp' ? 'up' : 'down';
+            if (!view.endOfTextblock(dir)) return false;
+
+            const isMermaid = (node: any) =>
+              node?.type === codeBlockType && node.attrs.language === 'mermaid';
+
+            if (event.key === 'ArrowUp') {
+              if ($from.parentOffset !== 0) return false;
+              const beforePos = $from.before();
+              const prev = state.doc.resolve(beforePos).nodeBefore;
+              if (!isMermaid(prev)) return false;
+              const endPos = beforePos - 1;
+              view.dispatch(
+                state.tr.setSelection(
+                  TextSelection.create(state.doc, endPos),
+                ),
+              );
+              return true;
+            }
+            if ($from.parentOffset !== $from.parent.nodeSize - 2) return false;
+            const afterPos = $from.after();
+            const next = state.doc.resolve(afterPos).nodeAfter;
+            if (!isMermaid(next)) return false;
+            const startPos = afterPos + 1;
+            view.dispatch(
+              state.tr.setSelection(
+                TextSelection.create(state.doc, startPos),
+              ),
+            );
+            return true;
+          },
+        },
+      }),
     ];
   },
 });