@marimo-team/islands 0.16.4 → 0.17.0

package/src/core/cells/cells.ts

@@ -5,7 +5,7 @@ import { type Atom, atom, useAtom, useAtomValue } from "jotai";
 import { atomFamily, selectAtom, splitAtom } from "jotai/utils";
 import { isEqual, zip } from "lodash-es";
 import { createRef, type ReducerWithoutAction } from "react";
-import type { CellHandle } from "@/components/editor/Cell";
+import type { CellHandle } from "@/components/editor/notebook-cell";
 import {
   type CellColumnId,
   type CellIndex,
@@ -1324,21 +1324,19 @@ const {
       cellRuntime: newCellRuntime,
     };
   },
-  upsertSetupCell: (state, action: { code: string }) => {
-    const { code } = action;
+  addSetupCellIfDoesntExist: (state, action: { code?: string }) => {
+    let { code } = action;
+    if (code == null) {
+      code = "# Initialization code that runs before all other cells";
+    }
 
     // First check if setup cell already exists
     if (SETUP_CELL_ID in state.cellData) {
-      // Update existing setup cell
-      return updateCellData({
-        state,
-        cellId: SETUP_CELL_ID,
-        cellReducer: (cell) => ({
-          ...cell,
-          code,
-          edited: code.trim() !== cell.lastCodeRun?.trim(),
-        }),
-      });
+      // Just focus on the existing setup cell
+      return {
+        ...state,
+        scrollKey: SETUP_CELL_ID,
+      };
     }
 
     return {
@@ -1365,6 +1363,7 @@ const {
         ...state.cellHandles,
         [SETUP_CELL_ID]: createRef(),
       },
+      scrollKey: SETUP_CELL_ID,
     };
   },
 });
@@ -1616,7 +1615,7 @@ export const columnIdsAtom = atom((get) =>
   get(notebookAtom).cellIds.getColumnIds(),
 );
 
-const cellDataAtom = atomFamily((cellId: CellId) =>
+export const cellDataAtom = atomFamily((cellId: CellId) =>
   atom((get) => get(notebookAtom).cellData[cellId]),
 );
 const cellRuntimeAtom = atomFamily((cellId: CellId) =>

package/src/core/cells/scrollCellIntoView.ts

@@ -1,10 +1,10 @@
 /* Copyright 2024 Marimo. All rights reserved. */
 import type { RefObject } from "react";
-import type { CellHandle } from "@/components/editor/Cell";
 import {
   isAnyCellFocused,
   tryFocus,
 } from "@/components/editor/navigation/focus-utils";
+import type { CellHandle } from "@/components/editor/notebook-cell";
 import { retryWithTimeout } from "@/utils/timeout";
 import { Logger } from "../../utils/Logger";
 import { goToVariableDefinition } from "../codemirror/go-to-definition/commands";

package/src/core/codemirror/__tests__/__snapshots__/setup.test.ts.snap

@@ -45,6 +45,7 @@ exports[`snapshot all duplicate keymaps > default keymaps 2`] = `
     },
     {
       "key": "Backspace",
+      "preventDefault": true,
       "run": "deleteCharBackward",
       "shift": "deleteCharBackward",
     },
@@ -153,6 +154,7 @@ exports[`snapshot all duplicate keymaps > vim keymaps 2`] = `
     },
     {
       "key": "Backspace",
+      "preventDefault": true,
       "run": "deleteCharBackward",
       "shift": "deleteCharBackward",
     },

package/src/core/codemirror/__tests__/replace-editor-content.test.ts

@@ -0,0 +1,336 @@
+/* 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 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({
+        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();
+  });
+
+  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();
+  });
+});

package/src/core/codemirror/__tests__/setup.test.ts

@@ -67,6 +67,7 @@ function getOpts() {
     hotkeys: new OverridingHotkeyProvider({}),
     theme: "light",
     displayConfig: { reference_highlighting: false },
+    inlineAiTooltip: false,
   } as const;
 }
 

package/src/core/codemirror/cm.ts

@@ -46,7 +46,6 @@ import type {
   KeymapConfig,
   LSPConfig,
 } from "../config/config-schema";
-import { getFeatureFlag } from "../config/feature-flag";
 import type { HotkeyProvider } from "../hotkeys/hotkeys";
 import { store } from "../state/jotai";
 import { requestEditCompletion } from "./ai/request";
@@ -83,6 +82,7 @@ export interface CodeMirrorSetupOpts {
   lspConfig: LSPConfig;
   diagnosticsConfig: DiagnosticsConfig;
   displayConfig: Pick<DisplayConfig, "reference_highlighting">;
+  inlineAiTooltip: boolean;
 }
 
 function getPlaceholderType(opts: CodeMirrorSetupOpts) {
@@ -104,6 +104,7 @@ export const setupCodeMirror = (opts: CodeMirrorSetupOpts): Extension[] => {
     lspConfig,
     diagnosticsConfig,
     displayConfig,
+    inlineAiTooltip,
   } = opts;
   const placeholderType = getPlaceholderType(opts);
 
@@ -128,7 +129,7 @@ export const setupCodeMirror = (opts: CodeMirrorSetupOpts): Extension[] => {
     goToDefinitionBundle(),
     diagnosticsConfig?.enabled ? lintGutter() : [],
     // AI edit inline
-    enableAI && getFeatureFlag("inline_ai_tooltip")
+    enableAI && inlineAiTooltip
       ? [
           aiExtension({
             prompt: (req) => {

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

package/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)],
   });
 }

package/src/core/codemirror/language/__tests__/sql.test.ts

@@ -8,6 +8,12 @@ import { PostgreSQL } from "@codemirror/lang-sql";
 import { EditorState, type Extension } from "@codemirror/state";
 import { DuckDBDialect } from "@marimo-team/codemirror-sql/dialects";
 import { afterEach, beforeEach, describe, expect, it } from "vitest";
+import type { CellId } from "@/core/cells/ids";
+import type {
+  CompletionConfig,
+  DiagnosticsConfig,
+  LSPConfig,
+} from "@/core/config/config-schema";
 import type { DataSourceConnection } from "@/core/datasets/data-source-connections";
 import {
   dataSourceConnectionsAtom,
@@ -16,7 +22,9 @@ import {
 import { type ConnectionName, DUCKDB_ENGINE } from "@/core/datasets/engines";
 import { datasetsAtom } from "@/core/datasets/state";
 import type { DatasetsState } from "@/core/datasets/types";
+import type { HotkeyProvider } from "@/core/hotkeys/hotkeys";
 import { store } from "@/core/state/jotai";
+import type { PlaceholderType } from "../../config/types";
 import { TestSQLCompletionStore } from "../languages/sql/completion-store";
 import {
   SQLLanguageAdapter,
@@ -28,6 +36,16 @@ const adapter = new SQLLanguageAdapter();
 
 const TEST_ENGINE = "test_engine" as ConnectionName;
 
+const TEST_EXTENSION_ARGS = [
+  {} as CellId,
+  {} as CompletionConfig,
+  {} as HotkeyProvider,
+  {} as PlaceholderType,
+  {} as LSPConfig & {
+    diagnostics: DiagnosticsConfig;
+  },
+] as const;
+
 describe("SQLLanguageAdapter", () => {
   describe("defaultMetadata", () => {
     it("should be set", () => {
@@ -1716,7 +1734,7 @@ describe("tablesCompletionSource", () => {
           const ctx = createCompletionContext(state, 14);
 
           const adapter = new SQLLanguageAdapter();
-          const extensions = adapter.getExtension();
+          const extensions = adapter.getExtension(...TEST_EXTENSION_ARGS);
           const completion = getCompletion(extensions);
 
           expect(completion).toBeDefined();
@@ -1779,7 +1797,7 @@ describe("tablesCompletionSource", () => {
           const ctx = createCompletionContext(state, 15, "u", 14);
 
           const adapter = new SQLLanguageAdapter();
-          const extensions = adapter.getExtension();
+          const extensions = adapter.getExtension(...TEST_EXTENSION_ARGS);
           const completion = getCompletion(extensions);
 
           expect(completion).toBeDefined();
@@ -1810,7 +1828,7 @@ describe("tablesCompletionSource", () => {
           const ctx = createCompletionContext(state, 15, "d", 14);
 
           const adapter = new SQLLanguageAdapter();
-          const extensions = adapter.getExtension();
+          const extensions = adapter.getExtension(...TEST_EXTENSION_ARGS);
           const completion = getCompletion(extensions);
 
           expect(completion).toBeDefined();
@@ -1843,7 +1861,7 @@ describe("tablesCompletionSource", () => {
           const ctx = createCompletionContext(state, 3, "SEL", 0);
 
           const adapter = new SQLLanguageAdapter();
-          const extensions = adapter.getExtension();
+          const extensions = adapter.getExtension(...TEST_EXTENSION_ARGS);
           const completion = getCompletion(extensions);
 
           expect(completion).toBeDefined();
@@ -1874,7 +1892,7 @@ describe("tablesCompletionSource", () => {
           const ctx = createCompletionContext(state, 14, ".n", 12);
 
           const adapter = new SQLLanguageAdapter();
-          const extensions = adapter.getExtension();
+          const extensions = adapter.getExtension(...TEST_EXTENSION_ARGS);
           const completion = getCompletion(extensions);
 
           expect(completion).toBeDefined();
@@ -1906,7 +1924,7 @@ describe("tablesCompletionSource", () => {
       describe("variableCompletionSource", () => {
         it("should be included in extension overrides", () => {
           const adapter = new SQLLanguageAdapter();
-          const extensions = adapter.getExtension();
+          const extensions = adapter.getExtension(...TEST_EXTENSION_ARGS);
           const completion = getCompletion(extensions);
 
           expect(completion).toBeDefined();

package/src/core/codemirror/language/languages/python.ts

@@ -128,7 +128,6 @@ const pylspClient = once((lspConfig: LSPConfig) => {
   return new NotebookLanguageServerClient(
     new LanguageServerClient({
       ...lspClientOpts,
-      autoClose: false,
     }),
     settings,
   );
@@ -146,7 +145,6 @@ const tyLspClient = once((_: LSPConfig) => {
   return new NotebookLanguageServerClient(
     new LanguageServerClient({
       ...lspClientOpts,
-      autoClose: false,
       getWorkspaceConfiguration: (_) => [{ disableLanguageServices: true }],
     }),
     {},
@@ -165,7 +163,6 @@ const pyrightClient = once((_: LSPConfig) => {
   return new NotebookLanguageServerClient(
     new LanguageServerClient({
       ...lspClientOpts,
-      autoClose: false,
     }),
     {},
   );
@@ -242,6 +239,7 @@ export class PythonLanguageAdapter implements LanguageAdapter<{}> {
             client: client as unknown as LanguageServerClient,
             languageId: "python",
             allowHTMLContent: true,
+            useSnippetOnCompletion: false,
             hoverConfig: hoverOptions,
             completionConfig: autocompleteOptions,
             // Default to false