@marimo-team/islands 0.25.3-dev170 → 0.25.3-dev175

package/src/components/editor/navigation/__tests__/navigation.test.ts

@@ -3,15 +3,37 @@
 
 import type { EditorView } from "@codemirror/view";
 import { act, renderHook } from "@testing-library/react";
-import { Provider } from "jotai";
 import React, { createRef } from "react";
-import { afterAll, beforeEach, describe, expect, it, vi } from "vitest";
+import {
+  afterAll,
+  afterEach,
+  beforeEach,
+  describe,
+  expect,
+  it,
+  vi,
+} from "vitest";
 import { Mocks } from "@/__mocks__/common";
+
+const { mockCreateCell, mockMoveToNextCell } = vi.hoisted(() => ({
+  mockCreateCell: vi.fn(),
+  mockMoveToNextCell: vi.fn(),
+}));
+
+import { DocumentWrapper, mockDocumentModel } from "@/__mocks__/document";
 import { MockNotebook } from "@/__mocks__/notebook";
 import { MockRequestClient } from "@/__mocks__/requests";
 import { aiCompletionCellAtom } from "@/core/ai/state";
-import type { CellActions } from "@/core/cells/cells";
+import type { CellActions } from "@/core/cells/actions";
 import { notebookAtom } from "@/core/cells/cells";
+import type { CellId } from "@/core/cells/ids";
+import type { DocumentModel } from "@/core/notebook/document";
+import {
+  cell,
+  CONFIG,
+  document as documentOf,
+  withCells,
+} from "@/core/notebook/document/__tests__/fixtures";
 import { LanguageAdapters } from "@/core/codemirror/language/LanguageAdapters";
 import { signatureHintField } from "@/core/codemirror/completion/signature-hint";
 import {
@@ -30,9 +52,11 @@ import {
 } from "../navigation";
 
 // Mock only the essential dependencies that we need to control
+vi.mock("@/core/cells/actions", () => ({
+  useCellActions: vi.fn(),
+}));
 vi.mock("@/core/cells/cells", async (importOriginal) => ({
   ...(await importOriginal()),
-  useCellActions: vi.fn(),
 }));
 
 vi.mock("@/core/cells/focus", async (importOriginal) => ({
@@ -80,7 +104,7 @@ vi.mock("@codemirror/autocomplete", () => ({
 
 // Get mocked functions
 const mockUseCellActions = vi.mocked(
-  await import("@/core/cells/cells"),
+  await import("@/core/cells/actions"),
 ).useCellActions;
 
 const mockUseSaveNotebook = vi.mocked(
@@ -93,6 +117,8 @@ const mockUseCellClipboard = vi.mocked(
   await import("../clipboard"),
 ).useCellClipboard;
 
+afterEach(() => stopModel());
+
 afterAll(() => {
   vi.resetAllMocks();
 });
@@ -107,11 +133,17 @@ import {
   useIsCellSelected,
 } from "../selection";
 
+// The document the hooks edit, with cellId3's code hidden
+let model: DocumentModel;
+let stopModel = () => {};
+const hideCodeOf = (id: CellId) =>
+  store.get(model.documentAtom)?.cells.get(id)?.config.hide_code;
+
 // Shared render helper
 const renderWithProvider = <T>(hook: () => T) => {
   return renderHook(hook, {
     wrapper: ({ children }) =>
-      React.createElement(Provider, { store }, children),
+      React.createElement(DocumentWrapper, { store, model }, children),
   });
 };
 
@@ -125,17 +157,16 @@ const mockCutCell = vi.fn().mockResolvedValue(undefined);
 const mockPasteCell = vi.fn();
 
 const mockCellActions = MockNotebook.cellActions({
+  createCell: mockCreateCell,
+  moveToNextCell: mockMoveToNextCell,
   focusCell: vi.fn(),
-  moveToNextCell: vi.fn(),
   focusTopCell: vi.fn(),
   focusBottomCell: vi.fn(),
-  createNewCell: vi.fn(),
   moveCell: vi.fn(),
   sendToTop: vi.fn(),
   sendToBottom: vi.fn(),
-  updateCellConfig: vi.fn(),
   markTouched: vi.fn(),
-  deleteCell: vi.fn(),
+  deleteCells: vi.fn(),
   undoDeleteCell: vi.fn(),
 });
 
@@ -207,6 +238,23 @@ describe("useCellNavigationProps", () => {
       },
     });
     store.set(notebookAtom, notebookState);
+    stopModel();
+    ({ model, stop: stopModel } = mockDocumentModel(
+      store,
+      withCells(
+        documentOf(),
+        {
+          c1: null,
+          c2: null,
+          [cellId1]: cell(cellId1, ""),
+          [cellId2]: cell(cellId2, ""),
+          [cellId3]: cell(cellId3, "", {
+            config: { ...CONFIG, hide_code: true },
+          }),
+        },
+        [cellId1, cellId2, cellId3],
+      ),
+    ));
 
     // Clear AI completion state
     store.set(aiCompletionCellAtom, null);
@@ -367,7 +415,7 @@ describe("useCellNavigationProps", () => {
       });
 
       expect(mockRunCell).toHaveBeenCalled();
-      expect(mockCellActions.moveToNextCell).toHaveBeenCalledWith({
+      expect(mockMoveToNextCell).toHaveBeenCalledWith({
         cellId: mockCellId,
         before: false,
       });
@@ -399,11 +447,7 @@ describe("useCellNavigationProps", () => {
         result.current.onKeyDown?.(mockEvent);
       });
 
-      expect(mockCellActions.createNewCell).toHaveBeenCalledWith({
-        cellId: mockCellId,
-        before: true,
-        autoFocus: true,
-      });
+      expect(mockCreateCell).toHaveBeenCalledWith({ before: mockCellId });
     });
 
     it("should create cell after when 'b' key is pressed", () => {
@@ -417,11 +461,7 @@ describe("useCellNavigationProps", () => {
         result.current.onKeyDown?.(mockEvent);
       });
 
-      expect(mockCellActions.createNewCell).toHaveBeenCalledWith({
-        cellId: mockCellId,
-        before: false,
-        autoFocus: true,
-      });
+      expect(mockCreateCell).toHaveBeenCalledWith({ after: mockCellId });
     });
 
     it("should add an import before and a SQL cell after when 'q' is pressed", () => {
@@ -435,18 +475,15 @@ describe("useCellNavigationProps", () => {
         result.current.onKeyDown?.(Mocks.keyboardEvent({ key: "q" }));
       });
 
-      expect(mockCellActions.createNewCell).toHaveBeenNthCalledWith(
+      expect(mockCreateCell).toHaveBeenNthCalledWith(
         1,
         expect.objectContaining({
-          cellId: mockCellId,
-          before: true,
+          before: mockCellId,
           code: "import marimo as mo",
         }),
       );
-      expect(mockCellActions.createNewCell).toHaveBeenNthCalledWith(2, {
-        cellId: mockCellId,
-        before: false,
-        autoFocus: true,
+      expect(mockCreateCell).toHaveBeenNthCalledWith(2, {
+        after: mockCellId,
         code: LanguageAdapters.sql.defaultCode,
       });
     });
@@ -1034,7 +1071,7 @@ describe("useCellNavigationProps", () => {
 
       expect(mockRunCell).toHaveBeenCalledWith([cellId1, cellId2]);
       // Should move to next cell after the last selected cell
-      expect(mockCellActions.moveToNextCell).toHaveBeenCalledWith({
+      expect(mockMoveToNextCell).toHaveBeenCalledWith({
         cellId: cellId2,
         before: false,
       });
@@ -1066,7 +1103,7 @@ describe("useCellNavigationProps", () => {
 
       expect(mockRunCell).toHaveBeenCalledWith([cellId2, cellId3]);
       // Should move to previous cell before the first selected cell
-      expect(mockCellActions.moveToNextCell).toHaveBeenCalledWith({
+      expect(mockMoveToNextCell).toHaveBeenCalledWith({
         cellId: cellId2,
         before: true,
       });
@@ -1167,15 +1204,7 @@ describe("useCellNavigationProps", () => {
         result.current.onKeyDown?.(mockEvent);
       });
 
-      expect(mockRequestClient.saveCellConfig).toHaveBeenCalledWith({
-        configs: {
-          [cellId1]: { hide_code: true },
-        },
-      });
-      expect(mockCellActions.updateCellConfig).toHaveBeenCalledWith({
-        cellId: cellId1,
-        config: { hide_code: true },
-      });
+      expect(hideCodeOf(cellId1)).toBe(true);
       expect(mockEvent.preventDefault).toHaveBeenCalled();
     });
 
@@ -1190,15 +1219,7 @@ describe("useCellNavigationProps", () => {
         result.current.onKeyDown?.(mockEvent);
       });
 
-      expect(mockRequestClient.saveCellConfig).toHaveBeenCalledWith({
-        configs: {
-          [cellId3]: { hide_code: false },
-        },
-      });
-      expect(mockCellActions.updateCellConfig).toHaveBeenCalledWith({
-        cellId: cellId3,
-        config: { hide_code: false },
-      });
+      expect(hideCodeOf(cellId3)).toBe(false);
     });
 
     it("should toggle hide code for multiple selected cells", () => {
@@ -1220,20 +1241,8 @@ describe("useCellNavigationProps", () => {
         result.current.onKeyDown?.(mockEvent);
       });
 
-      expect(mockRequestClient.saveCellConfig).toHaveBeenCalledWith({
-        configs: {
-          [cellId1]: { hide_code: true },
-          [cellId2]: { hide_code: true },
-        },
-      });
-      expect(mockCellActions.updateCellConfig).toHaveBeenCalledWith({
-        cellId: cellId1,
-        config: { hide_code: true },
-      });
-      expect(mockCellActions.updateCellConfig).toHaveBeenCalledWith({
-        cellId: cellId2,
-        config: { hide_code: true },
-      });
+      expect(hideCodeOf(cellId1)).toBe(true);
+      expect(hideCodeOf(cellId2)).toBe(true);
     });
   });
 
@@ -1701,7 +1710,7 @@ describe("useCellNavigationProps", () => {
         result.current.onKeyDown?.(mockEvent);
       });
 
-      expect(mockCellActions.createNewCell).not.toHaveBeenCalled();
+      expect(mockCreateCell).not.toHaveBeenCalled();
       expect(mockEvent.preventDefault).not.toHaveBeenCalled();
     });
 
@@ -1716,7 +1725,7 @@ describe("useCellNavigationProps", () => {
         result.current.onKeyDown?.(mockEvent);
       });
 
-      expect(mockCellActions.createNewCell).not.toHaveBeenCalled();
+      expect(mockCreateCell).not.toHaveBeenCalled();
       expect(mockEvent.preventDefault).not.toHaveBeenCalled();
     });
   });

package/src/components/editor/navigation/clipboard.ts

@@ -3,7 +3,8 @@
 import useEvent from "react-use-event-hook";
 import { z } from "zod";
 import { toast } from "@/components/ui/use-toast";
-import { getNotebook, useCellActions } from "@/core/cells/cells";
+import { getNotebook } from "@/core/cells/cells";
+import { useCellActions } from "@/core/cells/actions";
 import type { CellId } from "@/core/cells/ids";
 import {
   usePendingCutActions,
@@ -73,6 +74,7 @@ async function writeCellsToClipboard(
 
 export function useCellClipboard() {
   const actions = useCellActions();
+  const { createCell } = actions;
   const pendingCutActions = usePendingCutActions();
   const pendingCutState = usePendingCutState();
 
@@ -178,14 +180,15 @@ export function useCellClipboard() {
             const currentCellId = cellId;
             const reversedCells = clipboardData.cells.toReversed();
             for (const cell of reversedCells) {
-              actions.createNewCell({
-                cellId: currentCellId,
-                before,
+              // The pasted cell joins the column it is pasted into.
+              const { column: _column, ...config } = cell.config ?? {};
+              createCell({
+                ...(before
+                  ? { before: currentCellId }
+                  : { after: currentCellId }),
                 code: cell.code,
                 name: cell.name,
-                config: cell.config,
-                hideCode: cell.config?.hide_code,
-                autoFocus: true,
+                config,
               });
             }
 
@@ -199,11 +202,9 @@ export function useCellClipboard() {
       // Fallback to plain text
       const text = await navigator.clipboard.readText();
       if (text.trim()) {
-        actions.createNewCell({
-          cellId,
-          before,
+        createCell({
+          ...(before ? { before: cellId } : { after: cellId }),
           code: text,
-          autoFocus: true,
         });
       } else {
         toastNothingToPaste();

package/src/components/editor/navigation/multi-cell-action-toolbar.tsx

@@ -34,14 +34,14 @@ import {
   getCellEditorView,
   hasOnlyOneCellAtom,
   notebookAtom,
-  useCellActions,
 } from "@/core/cells/cells";
+import { useCellActions } from "@/core/cells/actions";
+import { useNotebookDocumentEdit } from "@/core/notebook/document";
 import type { CellId } from "@/core/cells/ids";
 import { usePendingDeleteService } from "@/core/cells/pending-delete-service";
 import { formatEditorViews } from "@/core/codemirror/format";
 import { userConfigAtom } from "@/core/config/config";
 import type { HotkeyAction } from "@/core/hotkeys/hotkeys";
-import { useRequestClient } from "@/core/network/requests";
 import type { CellConfig } from "@/core/network/types";
 import { store } from "@/core/state/jotai";
 import { useEventListener } from "@/hooks/useEventListener";
@@ -112,20 +112,15 @@ const CellStateDropdown: React.FC<{
 };
 
 export function useMultiCellActionButtons(cellIds: CellId[]) {
-  const {
-    updateCellConfig,
-    moveCell,
-    clearCellOutput,
-    sendToTop,
-    sendToBottom,
-  } = useCellActions();
+  const { moveCell, clearCellOutput, sendToTop, sendToBottom } =
+    useCellActions();
+  const edit = useNotebookDocumentEdit();
   const deleteCell = useDeleteManyCellsCallback();
   const hasOnlyOneCell = useAtomValue(hasOnlyOneCellAtom);
   const selectionActions = useCellSelectionActions();
   const runCells = useRunCells();
   const pendingDeleteService = usePendingDeleteService();
   const userConfig = useAtomValue(userConfigAtom);
-  const { saveCellConfig } = useRequestClient();
 
   const selectedCount = cellIds.length;
 
@@ -205,16 +200,12 @@ export function useMultiCellActionButtons(cellIds: CellId[]) {
   });
 
   const toggleSelectedCellsProperty = useEvent(
-    async (cellIds: CellId[], property: keyof CellConfig, value: boolean) => {
-      const configs: Record<CellId, Partial<CellConfig>> = {};
-      cellIds.forEach((cellId) => {
-        configs[cellId] = { [property]: value };
-      });
-
-      await saveCellConfig({ configs });
-
-      cellIds.forEach((cellId) => {
-        updateCellConfig({ cellId, config: configs[cellId] });
+    (cellIds: CellId[], property: keyof CellConfig, value: boolean) => {
+      edit((tx) => {
+        for (const cellId of cellIds) {
+          const config: Partial<CellConfig> = { [property]: value };
+          tx.updateConfig(cellId, config);
+        }
       });
     },
   );

package/src/components/editor/navigation/navigation.ts

@@ -13,7 +13,9 @@ import { mergeProps, useFocusWithin, useKeyboard } from "react-aria";
 import { DATA_FOR_CELL_ID } from "@/components/data-table/cell-utils";
 import { aiCompletionCellAtom } from "@/core/ai/state";
 import { maybeAddMarimoImport } from "@/core/cells/add-missing-import";
-import { cellIdsAtom, notebookAtom, useCellActions } from "@/core/cells/cells";
+import { cellIdsAtom, notebookAtom } from "@/core/cells/cells";
+import { useCellActions } from "@/core/cells/actions";
+import { useNotebookDocumentEdit } from "@/core/notebook/document";
 import { useCellFocusActions } from "@/core/cells/focus";
 import type { CellId } from "@/core/cells/ids";
 import { HTMLCellId } from "@/core/cells/ids";
@@ -36,7 +38,6 @@ import {
 } from "@/core/config/config";
 import type { HotkeyAction } from "@/core/hotkeys/hotkeys";
 import { parseShortcut } from "@/core/hotkeys/shortcuts";
-import { useRequestClient } from "@/core/network/requests";
 import { useSaveNotebook } from "@/core/saving/save-component";
 import { Events } from "@/utils/events";
 import type { CollapsibleTree } from "@/utils/id-tree";
@@ -189,9 +190,10 @@ export function useCellNavigationProps(
   },
 ) {
   const { saveOrNameNotebook } = useSaveNotebook();
-  const { saveCellConfig } = useRequestClient();
+  const edit = useNotebookDocumentEdit();
   const setAiCompletionCell = useSetAtom(aiCompletionCellAtom);
   const actions = useCellActions();
+  const { createCell, moveToNextCell } = useCellActions();
   const store = useStore();
   const temporarilyShownCodeActions = useTemporarilyShownCodeActions();
   const runCells = useRunCells();
@@ -376,21 +378,21 @@ export function useCellNavigationProps(
         "cell.runAndNewBelow": addSingleHandler((cellIds) => {
           runCells(cellIds);
           const lastCellId = cellIds[cellIds.length - 1];
-          actions.moveToNextCell({ cellId: lastCellId, before: false });
+          moveToNextCell({ cellId: lastCellId, before: false });
           return true;
         }),
         "cell.runAndNewAbove": addSingleHandler((cellIds) => {
           runCells(cellIds);
           const firstCellId = cellIds[0];
-          actions.moveToNextCell({ cellId: firstCellId, before: true });
+          moveToNextCell({ cellId: firstCellId, before: true });
           return true;
         }),
         "cell.createAbove": (cellId) => {
-          actions.createNewCell({ cellId, before: true });
+          createCell({ before: cellId });
           return true;
         },
         "cell.createBelow": (cellId) => {
-          actions.createNewCell({ cellId, before: false });
+          createCell({ after: cellId });
           return true;
         },
         "cell.moveUp": addSingleHandler((cellIds) => {
@@ -453,20 +455,12 @@ export function useCellNavigationProps(
             (config) => config?.hide_code,
           );
 
-          // Fire-and-forget
-          void saveCellConfig({
-            configs: Object.fromEntries(
-              cellIds.map((cellId) => [cellId, { hide_code: nextHideCode }]),
-            ),
+          edit((tx) => {
+            for (const cellId of cellIds) {
+              tx.updateConfig(cellId, { hide_code: nextHideCode });
+            }
           });
 
-          for (const cellId of cellIds) {
-            actions.updateCellConfig({
-              cellId,
-              config: { hide_code: nextHideCode },
-            });
-          }
-
           // Only focus if it is a single cell
           if (cellIds.length === 1) {
             const cellId = cellIds[0];
@@ -543,14 +537,14 @@ export function useCellNavigationProps(
           if (Events.hasModifier(evt)) {
             return false;
           }
-          actions.createNewCell({ cellId, before: true, autoFocus: true });
+          createCell({ before: cellId });
           return true;
         },
         "command.createCellAfter": (cellId) => {
           if (Events.hasModifier(evt)) {
             return false;
           }
-          actions.createNewCell({ cellId, before: false, autoFocus: true });
+          createCell({ after: cellId });
           return true;
         },
         "command.createSqlCellAfter": (cellId) => {
@@ -559,14 +553,12 @@ export function useCellNavigationProps(
           }
           maybeAddMarimoImport({
             autoInstantiate: true,
-            createNewCell: actions.createNewCell,
+            createCell: createCell,
             fromCellId: cellId,
             before: true,
           });
-          actions.createNewCell({
-            cellId,
-            before: false,
-            autoFocus: true,
+          createCell({
+            after: cellId,
             code: LanguageAdapters.sql.defaultCode,
           });
           return true;
@@ -635,11 +627,11 @@ export function useCellNavigationProps(
             return true;
           },
           o: () => {
-            actions.createNewCell({ cellId, before: false, autoFocus: true });
+            createCell({ after: cellId });
             return true;
           },
           "shift+o": () => {
-            actions.createNewCell({ cellId, before: true, autoFocus: true });
+            createCell({ before: cellId });
             return true;
           },
           u: () => {

package/src/components/editor/notebook-cell.tsx

@@ -43,13 +43,12 @@ import { useResizeObserver } from "@/hooks/useResizeObserver";
 import { cn } from "@/utils/cn";
 import type { Milliseconds, Seconds } from "@/utils/time";
 import {
-  type CellActions,
   createUntouchedCellAtom,
-  useCellActions,
   useCellData,
   useCellHandle,
   useCellRuntime,
 } from "../../core/cells/cells";
+import { type CellActions, useCellActions } from "../../core/cells/actions";
 import { type CellId, SETUP_CELL_ID } from "../../core/cells/ids";
 import {
   cellNeedsRun,
@@ -216,13 +215,10 @@ function useCellHiddenLogic({
 export type CellComponentActions = Pick<
   CellActions,
   | "updateCellCode"
-  | "createNewCell"
   | "focusCell"
   | "moveCell"
   | "collapseCell"
   | "expandCell"
-  | "moveToNextCell"
-  | "updateCellConfig"
   | "clearSerializedEditorState"
   | "setStdinResponse"
   | "clearCellConsoleOutput"
@@ -928,15 +924,16 @@ const CellLeftSideActions = memo(
     actions: CellComponentActions;
   }) => {
     const connection = useAtomValue(connectionAtom);
-    const { className, actions, cellId } = props;
+    const { className, cellId } = props;
+    const { createCell } = useCellActions();
 
     const createBelow = useEvent(
-      (opts: { code?: string; hideCode?: boolean } = {}) =>
-        actions.createNewCell({ cellId, before: false, ...opts }),
+      ({ code, hideCode }: { code?: string; hideCode?: boolean } = {}) =>
+        createCell({ after: cellId, code, config: { hide_code: hideCode } }),
     );
     const createAbove = useEvent(
-      (opts: { code?: string; hideCode?: boolean } = {}) =>
-        actions.createNewCell({ cellId, before: true, ...opts }),
+      ({ code, hideCode }: { code?: string; hideCode?: boolean } = {}) =>
+        createCell({ before: cellId, code, config: { hide_code: hideCode } }),
     );
 
     const oneClickShortcut = "mod";