From 4eefd697fa7fd37e12a06129f69a721f7c916514 Mon Sep 17 00:00:00 2001 From: Myles Scolnick Date: Wed, 15 Oct 2025 10:26:55 -0400 Subject: [PATCH 1/2] fix: preserve cursor position when updating cell code in watch mode or auto-formatting Fixes #6788 When using `marimo edit --watch` or auto-formatting, the cursor would jump to the top of a cell on save (or auto-save). This occurred because file watching would detect changes, send an `update-cell-codes` message to the frontend, and the frontend would replace the entire editor content without preserving the cursor position. This also fixes an old issue where the cursor would jump to the top of a cell after auto-formatting (regardless of whether it was on watch mode or not). ## Solution Created a new `replaceEditorContent()` utility that: - Checks if the editor has focus before replacing content - When focused, uses **line-based cursor preservation**: - Keeps the cursor on the same line number - Maintains the same column position within that line - If the line shrinks, clamps the cursor to the end of the line - If the line no longer exists, moves the cursor to the end of the last line - Prevents the cursor from jumping during external updates (like from file watching or auto-formatting) --- .../__tests__/replace-editor-content.test.ts | 269 ++++++++++++++++++ .../core/codemirror/find-replace/navigate.ts | 4 +- frontend/src/core/codemirror/format.ts | 11 +- .../src/core/codemirror/language/utils.ts | 11 +- .../core/codemirror/replace-editor-content.ts | 87 ++++++ frontend/src/core/saving/save-component.tsx | 79 ++--- 6 files changed, 406 insertions(+), 55 deletions(-) create mode 100644 frontend/src/core/codemirror/__tests__/replace-editor-content.test.ts create mode 100644 frontend/src/core/codemirror/replace-editor-content.ts diff --git a/frontend/src/core/codemirror/__tests__/replace-editor-content.test.ts b/frontend/src/core/codemirror/__tests__/replace-editor-content.test.ts new file mode 100644 index 00000000000..4817c4fbbf1 --- /dev/null +++ b/frontend/src/core/codemirror/__tests__/replace-editor-content.test.ts @@ -0,0 +1,269 @@ +/* Copyright 2024 Marimo. All rights reserved. */ + +import { EditorState } from "@codemirror/state"; +import { EditorView } from "@codemirror/view"; +import { describe, expect, it, vi } from "vitest"; +import { replaceEditorContent } from "../replace-editor-content"; + +describe("replaceEditorContent", () => { + it("should replace content when editor doesn't have focus", () => { + const view = new EditorView({ + state: EditorState.create({ + doc: "original content", + }), + }); + + // Editor doesn't have focus by default + expect(view.hasFocus).toBe(false); + + replaceEditorContent(view, "new content"); + + expect(view.state.doc.toString()).toBe("new content"); + // Cursor position is not preserved when not focused + expect(view.state.selection.main.head).toBe(0); + + view.destroy(); + }); + + it("should preserve cursor position when editor has focus (same line)", () => { + const view = new EditorView({ + state: EditorState.create({ + doc: "Hello World", + selection: { anchor: 6 }, // Position after "Hello " + }), + }); + + // Mock hasFocus to return true + Object.defineProperty(view, "hasFocus", { + get: () => true, + configurable: true, + }); + + expect(view.hasFocus).toBe(true); + + // Replace with similar length content on same line + replaceEditorContent(view, "Goodbye Everyone"); + + expect(view.state.doc.toString()).toBe("Goodbye Everyone"); + + // Cursor should stay at the same column (6) since it's still within the line + const newCursorPos = view.state.selection.main.head; + expect(newCursorPos).toBe(6); + + view.destroy(); + }); + + it("should preserve cursor at beginning when focused", () => { + const view = new EditorView({ + state: EditorState.create({ + doc: "Hello World", + selection: { anchor: 0 }, // At beginning + }), + }); + + Object.defineProperty(view, "hasFocus", { + get: () => true, + configurable: true, + }); + + replaceEditorContent(view, "Goodbye Everyone"); + + expect(view.state.doc.toString()).toBe("Goodbye Everyone"); + // Cursor should stay at beginning + expect(view.state.selection.main.head).toBe(0); + + view.destroy(); + }); + + it("should clamp cursor when line shrinks", () => { + const view = new EditorView({ + state: EditorState.create({ + doc: "Hello World", + selection: { anchor: 11 }, // At end (column 11) + }), + }); + + Object.defineProperty(view, "hasFocus", { + get: () => true, + configurable: true, + }); + + // Replace with shorter content + replaceEditorContent(view, "Goodbye"); + + expect(view.state.doc.toString()).toBe("Goodbye"); + // Cursor should be clamped to end of line since column 11 > line length (7) + expect(view.state.selection.main.head).toBe(7); + + view.destroy(); + }); + + it("should handle empty document", () => { + const view = new EditorView({ + state: EditorState.create({ + doc: "", + }), + }); + + Object.defineProperty(view, "hasFocus", { + get: () => true, + configurable: true, + }); + + replaceEditorContent(view, "new content"); + + expect(view.state.doc.toString()).toBe("new content"); + expect(view.state.selection.main.head).toBe(0); + + view.destroy(); + }); + + it("should do nothing when content is the same", () => { + const view = new EditorView({ + state: EditorState.create({ + doc: "same content", + selection: { anchor: 5 }, + }), + }); + + const dispatchSpy = vi.spyOn(view, "dispatch"); + + replaceEditorContent(view, "same content"); + + // No dispatch should have been called + expect(dispatchSpy).not.toHaveBeenCalled(); + expect(view.state.doc.toString()).toBe("same content"); + expect(view.state.selection.main.head).toBe(5); + + view.destroy(); + }); + + it("should handle cursor in middle of focused document", () => { + const view = new EditorView({ + state: EditorState.create({ + doc: "The quick brown fox jumps", + selection: { anchor: 10 }, // After "The quick " (column 10) + }), + }); + + Object.defineProperty(view, "hasFocus", { + get: () => true, + configurable: true, + }); + + // Replace with longer content + replaceEditorContent(view, "The extremely quick brown fox jumps over"); + + expect(view.state.doc.toString()).toBe( + "The extremely quick brown fox jumps over", + ); + + // Cursor should stay at same column (10) on same line + const newCursorPos = view.state.selection.main.head; + expect(newCursorPos).toBe(10); + + view.destroy(); + }); + + it("should respect preserveCursor=false when focused", () => { + const view = new EditorView({ + state: EditorState.create({ + doc: "Hello World", + selection: { anchor: 6 }, + }), + }); + + Object.defineProperty(view, "hasFocus", { + get: () => true, + configurable: true, + }); + + replaceEditorContent(view, "Goodbye Everyone", { preserveCursor: false }); + + expect(view.state.doc.toString()).toBe("Goodbye Everyone"); + // When preserveCursor is false, cursor is not explicitly set + // so it defaults to 0 + expect(view.state.selection.main.head).toBe(0); + + view.destroy(); + }); + + it("should handle newlines and multiline content", () => { + const view = new EditorView({ + state: EditorState.create({ + doc: "Line 1\nLine 2\nLine 3", + // Cursor at position 10: "Line 1\nLi|ne 2\nLine 3" + // Line 2, column 2 (after "Li") + selection: { anchor: 10 }, + }), + }); + + Object.defineProperty(view, "hasFocus", { + get: () => true, + configurable: true, + }); + + replaceEditorContent(view, "Line 1\nLine 2 updated\nLine 3\nLine 4"); + + expect(view.state.doc.toString()).toBe( + "Line 1\nLine 2 updated\nLine 3\nLine 4", + ); + + // Cursor should stay on line 2 at column 2 (after "Li") + // "Line 1\nLi|ne 2 updated\nLine 3\nLine 4" + const newCursorPos = view.state.selection.main.head; + expect(newCursorPos).toBe(10); // Same position, line 2 column 2 + + view.destroy(); + }); + + it("should move cursor up when line is deleted", () => { + const view = new EditorView({ + state: EditorState.create({ + doc: "Line 1\nLine 2\nLine 3", + selection: { anchor: 14 }, // Line 3, start of line + }), + }); + + Object.defineProperty(view, "hasFocus", { + get: () => true, + configurable: true, + }); + + // Replace with only 2 lines (line 3 is deleted) + replaceEditorContent(view, "Line 1\nLine 2"); + + expect(view.state.doc.toString()).toBe("Line 1\nLine 2"); + + // Cursor should move to end of last available line + const newCursorPos = view.state.selection.main.head; + expect(newCursorPos).toBe(13); // End of "Line 1\nLine 2" + + view.destroy(); + }); + + it("should preserve cursor on same line with column clamping", () => { + const view = new EditorView({ + state: EditorState.create({ + doc: "def function_with_long_name():", + selection: { anchor: 25 }, // Near end of line + }), + }); + + Object.defineProperty(view, "hasFocus", { + get: () => true, + configurable: true, + }); + + // Replace with shorter line + replaceEditorContent(view, "def fn():"); + + expect(view.state.doc.toString()).toBe("def fn():"); + + // Cursor should be clamped to end of shorter line + const newCursorPos = view.state.selection.main.head; + expect(newCursorPos).toBe(9); // End of "def fn():" + + view.destroy(); + }); +}); diff --git a/frontend/src/core/codemirror/find-replace/navigate.ts b/frontend/src/core/codemirror/find-replace/navigate.ts index 3c940c9d68d..6c67aba7a69 100644 --- a/frontend/src/core/codemirror/find-replace/navigate.ts +++ b/frontend/src/core/codemirror/find-replace/navigate.ts @@ -4,6 +4,7 @@ import { SearchQuery } from "@codemirror/search"; import { EditorSelection } from "@codemirror/state"; import { EditorView } from "@codemirror/view"; import { getAllEditorViews } from "@/core/cells/cells"; +import { replaceEditorContent } from "@/core/codemirror/replace-editor-content"; import { store } from "@/core/state/jotai"; import { asQueryCreator, type QueryType } from "./query"; import { findReplaceAtom } from "./state"; @@ -127,8 +128,7 @@ export const replaceAll = searchCommand(({ query }) => { const prevDoc = view.state.doc.toString(); undoHandlers.push(() => { - view.dispatch({ - changes: [{ from: 0, to: view.state.doc.length, insert: prevDoc }], + replaceEditorContent(view, prevDoc, { userEvent: "input.replace.all", }); }); diff --git a/frontend/src/core/codemirror/format.ts b/frontend/src/core/codemirror/format.ts index 4d9d1692ae3..661ccf3f1ba 100644 --- a/frontend/src/core/codemirror/format.ts +++ b/frontend/src/core/codemirror/format.ts @@ -17,6 +17,7 @@ import { getEditorCodeAsPython, updateEditorCodeFromPython, } from "./language/utils"; +import { replaceEditorContent } from "./replace-editor-content"; export const formattingChangeEffect = StateEffect.define(); @@ -106,15 +107,7 @@ export async function formatSQL(editor: EditorView) { }); // Update editor with formatted SQL - const doc = editor.state.doc; - - // Noop if the code is the same - if (doc.toString() === formattedSQL) { - return; - } - - editor.dispatch({ - changes: { from: 0, to: doc.length, insert: formattedSQL }, + replaceEditorContent(editor, formattedSQL, { effects: [formattingChangeEffect.of(true)], }); } diff --git a/frontend/src/core/codemirror/language/utils.ts b/frontend/src/core/codemirror/language/utils.ts index 8e9145df41b..72ac6d0678b 100644 --- a/frontend/src/core/codemirror/language/utils.ts +++ b/frontend/src/core/codemirror/language/utils.ts @@ -2,6 +2,7 @@ import type { EditorState } from "@codemirror/state"; import type { EditorView } from "@codemirror/view"; +import { replaceEditorContent } from "../replace-editor-content"; import { languageAdapterState } from "./extension"; import { languageMetadataField } from "./metadata"; @@ -36,14 +37,8 @@ export function updateEditorCodeFromPython( ): string { const languageAdapter = editor.state.field(languageAdapterState); const [code] = languageAdapter.transformIn(pythonCode); - const doc = editor.state.doc; - // Noop if the code is the same - if (doc.toString() === code) { - return code; - } - editor.dispatch({ - changes: { from: 0, to: doc.length, insert: code }, - }); + // Use replaceEditorContent which preserves cursor position when focused + replaceEditorContent(editor, code); return code; } diff --git a/frontend/src/core/codemirror/replace-editor-content.ts b/frontend/src/core/codemirror/replace-editor-content.ts new file mode 100644 index 00000000000..cf4b598e935 --- /dev/null +++ b/frontend/src/core/codemirror/replace-editor-content.ts @@ -0,0 +1,87 @@ +/* Copyright 2024 Marimo. All rights reserved. */ + +import type { TransactionSpec } from "@codemirror/state"; +import type { EditorView } from "@codemirror/view"; + +/** + * Replace the entire content of the editor with new content. + * + * When the editor has focus, this function attempts to preserve the cursor position + * by keeping it on the same line and column. If the line no longer exists or shrinks, + * the cursor is adjusted accordingly. This prevents the cursor from jumping during + * external updates (e.g., auto-formatting or external edits). + * + * When the editor doesn't have focus, the content is replaced without any special handling. + * + * @param editor - The CodeMirror editor view + * @param newContent - The new content to replace the editor with + * @param options - Optional configuration + * @param options.preserveCursor - Whether to preserve cursor position when focused (default: true) + * @param options.effects - Additional effects to apply to the transaction + * @param options.userEvent - User event string for the transaction + */ +export function replaceEditorContent( + editor: EditorView, + newContent: string, + options: { + preserveCursor?: boolean; + effects?: TransactionSpec["effects"]; + userEvent?: string; + } = {}, +): void { + const { preserveCursor = true, effects, userEvent } = options; + const doc = editor.state.doc; + const currentContent = doc.toString(); + + // Noop if content is the same + if (currentContent === newContent) { + return; + } + + // If editor has focus and we want to preserve cursor, try to maintain position + if (preserveCursor && editor.hasFocus) { + const cursorPos = editor.state.selection.main.head; + + // Get the current line and column + const currentLine = doc.lineAt(cursorPos); + const lineNumber = currentLine.number; // 1-based line number + const columnInLine = cursorPos - currentLine.from; // 0-based column in line + + // Calculate new cursor position based on line-aware logic + const newDoc = editor.state.update({ + changes: { from: 0, to: doc.length, insert: newContent }, + }).state.doc; + + let newCursorPos: number; + + if (lineNumber <= newDoc.lines) { + // Line still exists in new document + const newLine = newDoc.line(lineNumber); + const newLineLength = newLine.length; + + // Keep same column, but clamp to line length if line shrank + newCursorPos = newLine.from + Math.min(columnInLine, newLineLength); + } else { + // Line no longer exists, move to the end of the last line + const lastLine = newDoc.line(newDoc.lines); + newCursorPos = lastLine.to; + } + + // Apply changes with preserved cursor position + editor.dispatch({ + changes: { from: 0, to: doc.length, insert: newContent }, + selection: { + anchor: newCursorPos, + }, + effects, + userEvent, + }); + } else { + // Editor doesn't have focus or we don't want to preserve cursor + editor.dispatch({ + changes: { from: 0, to: doc.length, insert: newContent }, + effects, + userEvent, + }); + } +} diff --git a/frontend/src/core/saving/save-component.tsx b/frontend/src/core/saving/save-component.tsx index 72cf6dfdf09..0fdb4699091 100644 --- a/frontend/src/core/saving/save-component.tsx +++ b/frontend/src/core/saving/save-component.tsx @@ -96,54 +96,61 @@ export function useSaveNotebook() { const store = useStore(); // Save the notebook with the given filename - const saveNotebook = useEvent((filename: string, userInitiated: boolean) => { - const notebook = getNotebook(); - const cells = notebookCells(notebook); - const cellIds = cells.map((cell) => cell.id); - const codes = cells.map((cell) => cell.code); - const cellNames = cells.map((cell) => cell.name); - const configs = getCellConfigs(notebook); - const connection = store.get(connectionAtom); - const autoSaveConfig = store.get(autoSaveConfigAtom); - const layout = store.get(layoutStateAtom); - const kioskMode = store.get(kioskModeAtom); + const saveNotebook = useEvent( + async (filename: string, userInitiated: boolean) => { + const connection = store.get(connectionAtom); + const autoSaveConfig = store.get(autoSaveConfigAtom); + const kioskMode = store.get(kioskModeAtom); - if (kioskMode) { - return; - } + if (kioskMode) { + return; + } - // Don't save if there are no cells - if (codes.length === 0) { - return; - } + // Don't save if we are not connected to a kernel + if (connection.state !== WebSocketState.OPEN) { + openAlert("Failed to save notebook: not connected to a kernel."); + return; + } - // Don't save if we are not connected to a kernel - if (connection.state !== WebSocketState.OPEN) { - openAlert("Failed to save notebook: not connected to a kernel."); - return; - } + Logger.log("saving to ", filename); - Logger.log("saving to ", filename); - sendSave({ - cellIds: cellIds, - codes, - names: cellNames, - filename, - configs, - layout: getSerializedLayout(), - persist: true, - }).then(() => { if (userInitiated && autoSaveConfig.format_on_save) { - formatAll(); + Logger.log("formatting notebook (onSave)"); + await formatAll(); + } + + // Grab the latest notebook state, after formatting + const notebook = getNotebook(); + const cells = notebookCells(notebook); + const cellIds = cells.map((cell) => cell.id); + const codes = cells.map((cell) => cell.code); + const cellNames = cells.map((cell) => cell.name); + const configs = getCellConfigs(notebook); + const layout = store.get(layoutStateAtom); + + // Don't save if there are no cells + if (codes.length === 0) { + return; } + + await sendSave({ + cellIds: cellIds, + codes, + names: cellNames, + filename, + configs, + layout: getSerializedLayout(), + persist: true, + }); + setLastSavedNotebook({ names: cellNames, codes, configs, layout, }); - }); - }); + }, + ); // Save the notebook with the current filename, only if the filename exists const saveIfNotebookIsPersistent = useEvent((userInitiated = false) => { From 794a2d5c7a2b09944397440cd289328fb4ac73a1 Mon Sep 17 00:00:00 2001 From: Shahmir Varqha Date: Wed, 15 Oct 2025 23:15:49 +0800 Subject: [PATCH 2/2] add tests for cursor position (#6795) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## 📝 Summary was curious on some of these behaviours, added tests ## 🔍 Description of Changes ## 📋 Checklist - [x] I have read the [contributor guidelines](https://github.com/marimo-team/marimo/blob/main/CONTRIBUTING.md). - [ ] For large changes, or changes that affect the public API: this change was discussed or approved through an issue, on [Discord](https://marimo.io/discord?ref=pr), or the community [discussions](https://github.com/marimo-team/marimo/discussions) (Please provide a link if applicable). - [x] I have added tests for the changes made. - [x] I have run the code and verified that it works as expected. --- .../__tests__/replace-editor-content.test.ts | 67 +++++++++++++++++++ 1 file changed, 67 insertions(+) diff --git a/frontend/src/core/codemirror/__tests__/replace-editor-content.test.ts b/frontend/src/core/codemirror/__tests__/replace-editor-content.test.ts index 4817c4fbbf1..c9f686cae9d 100644 --- a/frontend/src/core/codemirror/__tests__/replace-editor-content.test.ts +++ b/frontend/src/core/codemirror/__tests__/replace-editor-content.test.ts @@ -242,6 +242,28 @@ describe("replaceEditorContent", () => { view.destroy(); }); + it("should stay at end of line when new line is added", () => { + const view = new EditorView({ + state: EditorState.create({ + doc: "Hello World", + selection: { anchor: 11 }, // At end (column 11) + }), + }); + + Object.defineProperty(view, "hasFocus", { + get: () => true, + configurable: true, + }); + + replaceEditorContent(view, "Hello World\nSome new line"); + + expect(view.state.doc.toString()).toBe("Hello World\nSome new line"); + const newCursorPos = view.state.selection.main.head; + expect(newCursorPos).toBe(11); + + view.destroy(); + }); + it("should preserve cursor on same line with column clamping", () => { const view = new EditorView({ state: EditorState.create({ @@ -266,4 +288,49 @@ describe("replaceEditorContent", () => { view.destroy(); }); + + it("should handle selection range (collapses to head position)", () => { + const view = new EditorView({ + state: EditorState.create({ + doc: "Hello World", + selection: { anchor: 0, head: 5 }, // "Hello" selected + }), + }); + + Object.defineProperty(view, "hasFocus", { + get: () => true, + configurable: true, + }); + + replaceEditorContent(view, "Goodbye Everyone"); + + expect(view.state.doc.toString()).toBe("Goodbye Everyone"); + // Selection head (5) is preserved as cursor position + expect(view.state.selection.main.head).toBe(5); + expect(view.state.selection.main.anchor).toBe(5); + + view.destroy(); + }); + + it("should handle replacing with empty string", () => { + const view = new EditorView({ + state: EditorState.create({ + doc: "Some content to clear", + selection: { anchor: 10 }, + }), + }); + + Object.defineProperty(view, "hasFocus", { + get: () => true, + configurable: true, + }); + + replaceEditorContent(view, ""); + + expect(view.state.doc.toString()).toBe(""); + // Cursor should be at position 0 (only valid position in empty doc) + expect(view.state.selection.main.head).toBe(0); + + view.destroy(); + }); });