fix: clear selection on undo/redo (#2385)

* fix: clear selection on undo/redo

* fix: don't drop prservedSelection for all transactions

* fix: address comments

Author: VladaHarbour

packages/super-editor/src/core/extensions/keymap-history.test.js

@@ -1,4 +1,5 @@
 import { describe, it, expect, afterEach } from 'vitest';
+import { TextSelection } from 'prosemirror-state';
 import { closeHistory, undoDepth } from 'prosemirror-history';
 import { initTestEditor } from '@tests/helpers/helpers.js';
 import { handleEnter, handleBackspace, handleDelete } from './keymap.js';
@@ -100,6 +101,51 @@ describe('keymap history grouping', () => {
     expect(editor.state.doc.textContent).toBe('hello');
   });
 
+  it('collapses selection after undo so layout does not treat it as active range', () => {
+    ({ editor } = initTestEditor({ mode: 'text', content: '<p>Hello world</p>' }));
+
+    // Select "Hello"
+    const from = 1;
+    const to = 6;
+    const sel = TextSelection.create(editor.state.doc, from, to);
+    editor.view.dispatch(editor.state.tr.setSelection(sel));
+
+    expect(editor.state.selection.from).toBe(from);
+    expect(editor.state.selection.to).toBe(to);
+    expect(editor.state.selection.empty).toBe(false);
+
+    // Simple edit to create an undo step
+    editor.view.dispatch(editor.state.tr.insertText('!', to));
+
+    // Undo should both revert the content change and collapse selection
+    editor.commands.undo();
+
+    const selectionAfterUndo = editor.state.selection;
+    expect(selectionAfterUndo.empty).toBe(true);
+  });
+
+  it('clears preservedSelection/lastSelection on undo so toolbar state does not resurrect old ranges', () => {
+    ({ editor } = initTestEditor({ mode: 'text', content: '<p>Hello world</p>' }));
+
+    // Seed editor-level selection snapshots (simulating toolbar/command preservation)
+    const from = 1;
+    const to = 6;
+    const sel = TextSelection.create(editor.state.doc, from, to);
+    editor.options.preservedSelection = sel;
+    editor.options.lastSelection = sel;
+
+    // Simple edit to create an undo step
+    editor.view.dispatch(editor.state.tr.insertText('!', to));
+
+    // Undo should trigger history cleanup, which clears editor-level selection snapshots
+    // and collapses any active text selection.
+    editor.commands.undo();
+
+    expect(editor.state.selection.empty).toBe(true);
+    expect(editor.options.preservedSelection).toBeNull();
+    expect(editor.options.lastSelection).toBeNull();
+  });
+
   it('closeHistory before deletion creates its own undo step', () => {
     ({ editor } = initTestEditor({ mode: 'text', content: '<p></p>' }));
 

packages/super-editor/src/extensions/custom-selection/custom-selection.js

@@ -5,7 +5,7 @@ import { Plugin, PluginKey, TextSelection } from 'prosemirror-state';
 import { Decoration, DecorationSet } from 'prosemirror-view';
 import { shouldAllowNativeContextMenu } from '../../utils/contextmenu-helpers.js';
 
-const DEFAULT_SELECTION_STATE = Object.freeze({
+export const DEFAULT_SELECTION_STATE = Object.freeze({
   focused: false,
   preservedSelection: null,
   showVisualSelection: false,
@@ -378,6 +378,14 @@ export const CustomSelection = Extension.create({
                   skipFocusReset: false,
                 }),
               );
+
+              // Also clear editor-level preserved selection snapshots so that
+              // subsequent commands (linked styles, mark commands, etc.) don't
+              // resurrect an old selection after history undo/redo.
+              this.editor.setOptions({
+                preservedSelection: null,
+                lastSelection: null,
+              });
             }
           },
         },

packages/super-editor/src/extensions/history/history.js

@@ -1,7 +1,49 @@
 // @ts-nocheck
+import { TextSelection } from 'prosemirror-state';
 import { history, redo as originalRedo, undo as originalUndo } from 'prosemirror-history';
 import { undo as yUndo, redo as yRedo, yUndoPlugin } from 'y-prosemirror';
 import { Extension } from '@core/Extension.js';
+import { CustomSelectionPluginKey, DEFAULT_SELECTION_STATE } from '../custom-selection/custom-selection.js';
+
+function applySelectionCleanup(editor, tr) {
+  let cleaned = tr.setMeta(CustomSelectionPluginKey, DEFAULT_SELECTION_STATE);
+
+  const sel = cleaned.selection;
+  if (sel && sel instanceof TextSelection && !sel.empty) {
+    try {
+      const collapsed = TextSelection.create(cleaned.doc, sel.head);
+      cleaned = cleaned.setSelection(collapsed);
+    } catch {
+      // Ignore collapse failures and fall back to original selection
+    }
+  }
+
+  editor.setOptions({
+    preservedSelection: null,
+    lastSelection: null,
+  });
+
+  return cleaned;
+}
+
+function createHistoryDispatch(editor, dispatch) {
+  if (!dispatch) return dispatch;
+  return (historyTr) => {
+    const cleaned = applySelectionCleanup(editor, historyTr);
+    dispatch(cleaned);
+  };
+}
+
+function runSelectionCleanupAfterCollabHistory(editor) {
+  const view = editor?.view;
+  const state = editor?.state;
+  if (!view || !state) return;
+
+  let tr = applySelectionCleanup(editor, state.tr);
+  // Avoid creating a new undo step for this synthetic cleanup transaction.
+  tr = tr.setMeta('addToHistory', false);
+  view.dispatch(tr);
+}
 
 /**
  * Configuration options for History
@@ -52,10 +94,13 @@ export const History = Extension.create({
       undo: () => ({ state, dispatch, tr }) => {
         if (this.editor.options.collaborationProvider && this.editor.options.ydoc) {
           tr.setMeta('preventDispatch', true);
-          return yUndo(state);
+          const result = yUndo(state);
+          runSelectionCleanupAfterCollabHistory(this.editor);
+          return result;
         }
         tr.setMeta('inputType', 'historyUndo');
-        return originalUndo(state, dispatch);
+        const wrappedDispatch = createHistoryDispatch(this.editor, dispatch);
+        return originalUndo(state, wrappedDispatch);
       },
 
       /**
@@ -68,10 +113,13 @@ export const History = Extension.create({
       redo: () => ({ state, dispatch, tr }) => {
         if (this.editor.options.collaborationProvider && this.editor.options.ydoc) {
           tr.setMeta('preventDispatch', true);
-          return yRedo(state);
+          const result = yRedo(state);
+          runSelectionCleanupAfterCollabHistory(this.editor);
+          return result;
         }
         tr.setMeta('inputType', 'historyRedo');
-        return originalRedo(state, dispatch);
+        const wrappedDispatch = createHistoryDispatch(this.editor, dispatch);
+        return originalRedo(state, wrappedDispatch);
       },
     };
   },