fix: preserve cursor position when updating cell code in watch mode or auto-formatting

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();
+  });
+});

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",
       });
     });

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<boolean>();
 
@@ -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)],
   });
 }

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;
 }