[lexical-code][lexical-playground] Bug Fix: Fix selection boundaries in code block (#7187)

packages/lexical-code/src/CodeHighlighter.ts

@@ -880,22 +880,64 @@ export function registerCodeHighlighting(
     ),
     editor.registerCommand(
       KEY_ARROW_UP_COMMAND,
-      (payload): boolean => $handleShiftLines(KEY_ARROW_UP_COMMAND, payload),
+      (event) => {
+        const selection = $getSelection();
+        if (!$isRangeSelection(selection)) {
+          return false;
+        }
+        const {anchor} = selection;
+        const anchorNode = anchor.getNode();
+        if (!$isSelectionInCode(selection)) {
+          return false;
+        }
+        // If at the start of a code block, prevent selection from moving out
+        if (
+          selection.isCollapsed() &&
+          anchor.offset === 0 &&
+          anchorNode.getPreviousSibling() === null &&
+          $isCodeNode(anchorNode.getParentOrThrow())
+        ) {
+          event.preventDefault();
+          return true;
+        }
+        return $handleShiftLines(KEY_ARROW_UP_COMMAND, event);
+      },
       COMMAND_PRIORITY_LOW,
     ),
     editor.registerCommand(
       KEY_ARROW_DOWN_COMMAND,
-      (payload): boolean => $handleShiftLines(KEY_ARROW_DOWN_COMMAND, payload),
+      (event) => {
-      COMMAND_PRIORITY_LOW,
+        const selection = $getSelection();
-    ),
+        if (!$isRangeSelection(selection)) {
-    editor.registerCommand(
+          return false;
-      MOVE_TO_END,
+        }
-      (payload): boolean => $handleMoveTo(MOVE_TO_END, payload),
+        const {anchor} = selection;
+        const anchorNode = anchor.getNode();
+        if (!$isSelectionInCode(selection)) {
+          return false;
+        }
+        // If at the end of a code block, prevent selection from moving out
+        if (
+          selection.isCollapsed() &&
+          anchor.offset === anchorNode.getTextContentSize() &&
+          anchorNode.getNextSibling() === null &&
+          $isCodeNode(anchorNode.getParentOrThrow())
+        ) {
+          event.preventDefault();
+          return true;
+        }
+        return $handleShiftLines(KEY_ARROW_DOWN_COMMAND, event);
+      },
       COMMAND_PRIORITY_LOW,
     ),
     editor.registerCommand(
       MOVE_TO_START,
-      (payload): boolean => $handleMoveTo(MOVE_TO_START, payload),
+      (event) => $handleMoveTo(MOVE_TO_START, event as KeyboardEvent),
+      COMMAND_PRIORITY_LOW,
+    ),
+    editor.registerCommand(
+      MOVE_TO_END,
+      (event) => $handleMoveTo(MOVE_TO_END, event as KeyboardEvent),
       COMMAND_PRIORITY_LOW,
     ),
   );

packages/lexical-playground/__tests__/e2e/CodeBlock.spec.mjs

@@ -17,6 +17,7 @@ import {
   assertHTML,
   assertSelection,
   click,
+  expect,
   focusEditor,
   html,
   initialize,
@@ -916,6 +917,162 @@ test.describe('CodeBlock', () => {
     await assertHTML(page, bcaHTML);
   });
 
+  test('prevents selection and typing outside code block boundaries', async ({
+    page,
+    isPlainText,
+  }) => {
+    test.skip(isPlainText);
+
+    await focusEditor(page);
+    await page.keyboard.type('console.log("test");');
+    await selectAll(page);
+    await toggleCodeBlock(page);
+
+    // Test 1: Selection stays at start when pressing up
+    await moveToStart(page);
+    await page.keyboard.press('ArrowUp');
+    await assertSelection(page, {
+      anchorOffset: 0,
+      anchorPath: [0, 0, 0],
+      focusOffset: 0,
+      focusPath: [0, 0, 0],
+    });
+
+    // Test 2: Typing at start stays within code block
+    await page.keyboard.type('// start');
+    await page.keyboard.press('Enter');
+    await assertHTML(
+      page,
+      html`
+        <code
+          class="PlaygroundEditorTheme__code PlaygroundEditorTheme__ltr"
+          dir="ltr"
+          spellcheck="false"
+          data-gutter="12"
+          data-highlight-language="javascript"
+          data-language="javascript">
+          <span
+            class="PlaygroundEditorTheme__tokenComment"
+            data-lexical-text="true">
+            // start
+          </span>
+          <br />
+          <span data-lexical-text="true">console</span>
+          <span
+            class="PlaygroundEditorTheme__tokenPunctuation"
+            data-lexical-text="true">
+            .
+          </span>
+          <span
+            class="PlaygroundEditorTheme__tokenFunction"
+            data-lexical-text="true">
+            log
+          </span>
+          <span
+            class="PlaygroundEditorTheme__tokenPunctuation"
+            data-lexical-text="true">
+            (
+          </span>
+          <span
+            class="PlaygroundEditorTheme__tokenSelector"
+            data-lexical-text="true">
+            "test"
+          </span>
+          <span
+            class="PlaygroundEditorTheme__tokenPunctuation"
+            data-lexical-text="true">
+            )
+          </span>
+          <span
+            class="PlaygroundEditorTheme__tokenPunctuation"
+            data-lexical-text="true">
+            ;
+          </span>
+        </code>
+      `,
+    );
+
+    // Let's verify the cursor position after typing the start comment
+    await assertSelection(page, {
+      anchorOffset: 0,
+      anchorPath: [0, 2, 0],
+      focusOffset: 0,
+      focusPath: [0, 2, 0],
+    });
+
+    // Test 3: Selection stays at end when pressing down
+    await moveToEnd(page);
+    await page.keyboard.type(' // end');
+    await assertHTML(
+      page,
+      html`
+        <code
+          class="PlaygroundEditorTheme__code PlaygroundEditorTheme__ltr"
+          dir="ltr"
+          spellcheck="false"
+          data-gutter="12"
+          data-highlight-language="javascript"
+          data-language="javascript">
+          <span
+            class="PlaygroundEditorTheme__tokenComment"
+            data-lexical-text="true">
+            // start
+          </span>
+          <br />
+          <span data-lexical-text="true">console</span>
+          <span
+            class="PlaygroundEditorTheme__tokenPunctuation"
+            data-lexical-text="true">
+            .
+          </span>
+          <span
+            class="PlaygroundEditorTheme__tokenFunction"
+            data-lexical-text="true">
+            log
+          </span>
+          <span
+            class="PlaygroundEditorTheme__tokenPunctuation"
+            data-lexical-text="true">
+            (
+          </span>
+          <span
+            class="PlaygroundEditorTheme__tokenSelector"
+            data-lexical-text="true">
+            "test"
+          </span>
+          <span
+            class="PlaygroundEditorTheme__tokenPunctuation"
+            data-lexical-text="true">
+            )
+          </span>
+          <span
+            class="PlaygroundEditorTheme__tokenPunctuation"
+            data-lexical-text="true">
+            ;
+          </span>
+          <span data-lexical-text="true"></span>
+          <span
+            class="PlaygroundEditorTheme__tokenComment"
+            data-lexical-text="true">
+            // end
+          </span>
+        </code>
+      `,
+    );
+
+    await page.keyboard.press('ArrowDown');
+    await assertSelection(page, {
+      anchorOffset: 6,
+      anchorPath: [0, 10, 0],
+      focusOffset: 6,
+      focusPath: [0, 10, 0],
+    });
+
+    // Verify no content escaped the code block
+    const paragraphs = await page.$$('p');
+    expect(paragraphs.length).toBe(0);
+  });
+
   test('When pressing CMD/Ctrl + Left, CMD/Ctrl + Right, the cursor should go to the start of the code', async ({
     page,
     isPlainText,