@marimo-team/islands 0.25.3-dev173 → 0.25.3-dev176

package/src/core/cells/add-missing-import.ts

@@ -3,7 +3,8 @@
 import { getRequestClient } from "../network/requests";
 import { store } from "../state/jotai";
 import { variablesAtom } from "../variables/state";
-import { type CellActions, notebookAtom } from "./cells";
+import { notebookAtom } from "./cells";
+import type { CreateCell, CreateCellOptions } from "./actions";
 import { CellId } from "./ids";
 
 /**
@@ -50,82 +51,85 @@ export function maybeAddMissingImport({
 /**
  * Adds a marimo import to the notebook if not already present.
  * @param autoInstantiate Whether to automatically run the cell.
- * @param createNewCell The function to create a new cell.
- * @param fromCellId The cell to add the import to.
+ * @param createCell The command that creates a cell.
+ * @param fromCellId The cell to add the import next to.
  * @param before Whether to add the import before or after the cell.
  *
  * Returns the ID of the new cell if added, otherwise null.
  */
 export function maybeAddMarimoImport({
   autoInstantiate,
-  createNewCell,
+  createCell,
   fromCellId,
   before,
 }: {
   autoInstantiate: boolean;
-  createNewCell: CellActions["createNewCell"];
+  createCell: CreateCell;
   fromCellId?: CellId | null;
   before?: boolean;
 }): CellId | null {
-  const client = getRequestClient();
-  let newCellId: CellId | null = null;
-  const added = maybeAddMissingImport({
+  return addImport({
     moduleName: "marimo",
     variableName: "mo",
-    onAddImport: (importStatement) => {
-      newCellId = CellId.create();
-      createNewCell({
-        cellId: fromCellId ?? "__end__",
-        before: before ?? false,
-        code: importStatement,
-        lastCodeRun: autoInstantiate ? importStatement : undefined,
-        newCellId: newCellId,
-        skipIfCodeExists: true,
-        autoFocus: false,
-      });
-      if (autoInstantiate) {
-        void client.sendRun({
-          cellIds: [newCellId],
-          codes: [importStatement],
-        });
-      }
-    },
+    autoInstantiate,
+    createCell,
+    placement: before ? { before: fromCellId } : { after: fromCellId },
   });
-  return added ? newCellId : null;
 }
 
 export function maybeAddAltairImport({
   autoInstantiate,
-  createNewCell,
+  createCell,
   fromCellId,
 }: {
   autoInstantiate: boolean;
-  createNewCell: CellActions["createNewCell"];
+  createCell: CreateCell;
   fromCellId?: CellId | null;
 }): CellId | null {
-  const client = getRequestClient();
-  let newCellId: CellId | null = null;
-  const added = maybeAddMissingImport({
+  return addImport({
     moduleName: "altair",
     variableName: "alt",
+    autoInstantiate,
+    createCell,
+    placement: { after: fromCellId },
+  });
+}
+
+/** Create the import cell, and run it when the notebook runs on its own. */
+function addImport({
+  moduleName,
+  variableName,
+  autoInstantiate,
+  createCell,
+  placement,
+}: {
+  moduleName: string;
+  variableName: string;
+  autoInstantiate: boolean;
+  createCell: CreateCell;
+  placement: Pick<CreateCellOptions, "after" | "before">;
+}): CellId | null {
+  let newCellId: CellId | null = null;
+  maybeAddMissingImport({
+    moduleName,
+    variableName,
     onAddImport: (importStatement) => {
-      newCellId = CellId.create();
-      createNewCell({
-        cellId: fromCellId ?? "__end__",
-        before: false,
+      const id = CellId.create();
+      const created = createCell({
+        ...placement,
+        id,
         code: importStatement,
-        lastCodeRun: autoInstantiate ? importStatement : undefined,
-        newCellId: newCellId,
         skipIfCodeExists: true,
-        autoFocus: false,
+        focus: false,
       });
-      if (autoInstantiate) {
-        void client.sendRun({
+      newCellId = created === null ? null : id;
+      if (autoInstantiate && newCellId !== null) {
+        void getRequestClient().sendRun({
           cellIds: [newCellId],
           codes: [importStatement],
         });
       }
     },
   });
-  return added ? newCellId : null;
+  return newCellId;
 }

package/src/core/cells/cells.ts

@@ -1,12 +1,11 @@
 /* Copyright 2026 Marimo. All rights reserved. */
 
-import { historyField } from "@codemirror/commands";
+import type * as Protocol from "@marimo-team/client/protocol";
 import { type Atom, atom, useAtom, useAtomValue } from "jotai";
 import { atomFamily, selectAtom, splitAtom } from "jotai/utils";
 import { createRef, type ReducerWithoutAction } from "react";
 import type { CellHandle } from "@/components/editor/notebook-cell";
-import type { EditableDocument } from "@/core/notebook/document";
-import type { CollapsibleTree } from "@/utils/id-tree";
+import type { EditableDocument, Placement } from "@/core/notebook/document";
 import {
   type CellColumnId,
   type CellIndex,
@@ -31,7 +30,7 @@ import { createDeepEqualAtom, store } from "../state/jotai";
 import { isWasm } from "../wasm/utils";
 import { prepareCellForExecution, transitionCell } from "./cell";
 import { applyDocument, documentEditMiddleware } from "./document";
-import { CellId, SCRATCH_CELL_ID, SETUP_CELL_ID } from "./ids";
+import { CellId, SCRATCH_CELL_ID } from "./ids";
 import { type CellLog, getCellLogsForMessage } from "./logs";
 import {
   focusAndScrollCellIntoView,
@@ -42,7 +41,6 @@ import {
   type CellData,
   type CellRuntimeState,
   createCell,
-  createCellConfig,
   createCellRuntimeState,
   isEdited,
 } from "./types";
@@ -56,17 +54,21 @@ import {
   notebookQueueOrRunningCount,
 } from "./utils";
 
+/** A deleted cell as it was, with where it went in the document. */
+export interface DeletedCell {
+  cell: Protocol.CellInput;
+  placement: Placement;
+  /** The editor's state at deletion, so undo brings its history back. */
+  serializedEditorState: SerializedEditorState | null;
+}
+
 /**
- * History entry for undoing a cell deletion.
+ * History entry for undoing one deletion, of one or more cells.
  */
 export interface UndoDeleteEntry {
   type: "delete";
-  name: string;
-  serializedEditorState: SerializedEditorState;
-  column: CellColumnId;
-  index: CellIndex;
-  isSetupCell: boolean;
-  config: CellConfig;
+  /** In document order. Each cell is placed after the one before it. */
+  cells: DeletedCell[];
 }
 
 /**
@@ -162,40 +164,6 @@ export function initialNotebookState(): NotebookState {
   });
 }
 
-/** The target cell ID to create a new cell relative to. Can be:
- * - A CellId string for an existing cell
- * - "__end__" to append at the end of the first column
- * - {type: "__end__", columnId} to append at the end of a specific column
- */
-export type CellPosition =
-  | CellId
-  | "__end__"
-  | { type: "__end__"; columnId: CellColumnId };
-
-export interface CreateNewCellAction {
-  cellId: CellPosition;
-  /** Whether to insert before (true) or after (false) the target cell */
-  before: boolean;
-  /** Initial code content for the new cell */
-  code?: string;
-  /** Optional name for the new cell */
-  name?: string;
-  /** Optional cell configuration */
-  config?: CellConfig;
-  /** The last executed code for the new cell */
-  lastCodeRun?: string;
-  /** Timestamp of the last execution */
-  lastExecutionTime?: number;
-  /** Optional custom ID for the new cell. Auto-generated if not provided */
-  newCellId?: CellId;
-  /** Whether to focus the new cell after creation */
-  autoFocus?: boolean;
-  /** If true, skip creation if code already exists */
-  skipIfCodeExists?: boolean;
-  /** Hide the code in the new cell. This will be initially shown until the cell is blurred for the first time. */
-  hideCode?: boolean;
-}
-
 /**
  * Actions and reducer for the notebook state.
  */
@@ -206,85 +174,6 @@ const {
   useActions,
   valueAtom: notebookAtom,
 } = createReducerAndAtoms(initialNotebookState, {
-  createNewCell: (state, action: CreateNewCellAction) => {
-    const {
-      cellId,
-      before,
-      code,
-      name,
-      config,
-      lastCodeRun = null,
-      lastExecutionTime = null,
-      autoFocus = true,
-      skipIfCodeExists = false,
-      hideCode = undefined,
-    } = action;
-
-    let columnId: CellColumnId;
-    let cellIndex: number;
-
-    // If skipIfCodeExists is true, check if the code already exists in the notebook
-    if (skipIfCodeExists) {
-      for (const cellId of state.cellIds.inOrderIds) {
-        if (state.cellData[cellId]?.code === code) {
-          return state;
-        }
-      }
-    }
-
-    if (cellId === "__end__") {
-      const column = state.cellIds.atOrThrow(0);
-      columnId = column.id;
-      cellIndex = column.length;
-    } else if (typeof cellId === "string") {
-      const column = state.cellIds.findWithId(cellId);
-      columnId = column.id;
-      cellIndex = column.topLevelIds.indexOf(cellId);
-    } else if (cellId.type === "__end__") {
-      const column =
-        state.cellIds.get(cellId.columnId) || state.cellIds.atOrThrow(0);
-      columnId = column.id;
-      cellIndex = column.length;
-    } else {
-      throw new Error("Invalid cellId");
-    }
-
-    const newCellId = action.newCellId || CellId.create();
-    const insertionIndex = before ? cellIndex : cellIndex + 1;
-
-    return {
-      ...state,
-      cellIds: state.cellIds.insertId(newCellId, columnId, insertionIndex),
-      cellData: {
-        ...state.cellData,
-        [newCellId]: createCell({
-          id: newCellId,
-          code,
-          name,
-          lastCodeRun,
-          config: createCellConfig({
-            ...config,
-            ...(hideCode != null && { hide_code: hideCode }),
-          }),
-          lastExecutionTime,
-          edited: Boolean(code) && code !== lastCodeRun,
-        }),
-      },
-      cellRuntime: {
-        ...state.cellRuntime,
-        [newCellId]: createCellRuntimeState(),
-      },
-      cellHandles: {
-        ...state.cellHandles,
-        [newCellId]: createRef(),
-      },
-      scrollKey: autoFocus ? newCellId : null,
-      untouchedNewCells:
-        hideCode && autoFocus
-          ? new Set([...state.untouchedNewCells, newCellId])
-          : state.untouchedNewCells,
-    };
-  },
   moveCell: (
     state,
     action: {
@@ -679,128 +568,29 @@ const {
       cellIds: state.cellIds.compact(),
     };
   },
-  deleteCell: (state, action: { cellId: CellId }) => {
-    const cellId = action.cellId;
-
-    // Can't delete the last cell, across all columns
-    if (state.cellIds.hasOnlyOneId()) {
-      return state;
-    }
-
-    let column: CollapsibleTree<CellId>;
-    try {
-      column = state.cellIds.findWithId(cellId);
-    } catch (error) {
-      // Expected for kernel-only cells or out-of-order transactions.
-      Logger.warn("Skipping delete for missing cellId", {
-        cellId,
-        error,
-      });
-      return state;
-    }
-
-    const cellIndex = column.indexOfOrThrow(cellId);
-    const focusIndex = cellIndex === 0 ? 1 : cellIndex - 1;
-    let scrollKey: CellId | null = null;
-    if (column.length > 1) {
-      scrollKey = column.atOrThrow(focusIndex);
-    }
-
-    const editorView = state.cellHandles[cellId].current?.editorView;
-    const serializedEditorState = editorView?.state.toJSON({
-      history: historyField,
-    });
-    if (serializedEditorState) {
-      serializedEditorState.doc = state.cellData[cellId].code;
-    }
-
-    // release the granular atom(s) created for this cell
-    releaseCellAtoms(cellId);
-    const prevData = state.cellData[cellId];
-    return {
-      ...state,
-      cellIds: state.cellIds.deleteById(cellId),
-      history: [
-        ...state.history,
-        {
-          type: "delete",
-          name: prevData.name,
-          serializedEditorState: serializedEditorState,
-          column: column.id,
-          index: cellIndex,
-          isSetupCell: cellId === SETUP_CELL_ID,
-          config: prevData.config,
-        },
-      ],
-      scrollKey: scrollKey,
-    };
-  },
-  undoDeleteCell: (state) => {
-    if (state.history.length === 0) {
+  /**
+   * Put moved cells back where they were.
+   *
+   * TODO(MO-8028): Move to `undoDeleteCellAtom` with the order family.
+   */
+  undoMoveCells: (state, action: UndoMoveEntry) => {
+    const { cellIds, placements } = action;
+    if (
+      cellIds.length === 0 ||
+      placements.length !== cellIds.length ||
+      cellIds.some((id) => !state.cellData[id])
+    ) {
       return state;
     }
-
-    const last = state.history[state.history.length - 1];
-
-    if (last.type === "move") {
-      const { cellIds, placements } = last;
-      if (
-        cellIds.length === 0 ||
-        placements.length !== cellIds.length ||
-        cellIds.some((id) => !state.cellData[id])
-      ) {
-        return { ...state, history: state.history.slice(0, -1) };
-      }
-      const toRestore = cellIds.map((id, i) => ({
-        id,
-        columnId: placements[i].columnId,
-        index: placements[i].index,
-      }));
-      return {
-        ...state,
-        cellIds: state.cellIds.placeCells(toRestore),
-        history: state.history.slice(0, -1),
-        scrollKey: cellIds[0] ?? null,
-      };
-    }
-
-    const {
-      name,
-      serializedEditorState = { doc: "" },
-      column,
-      index,
-      isSetupCell,
-      config,
-    } = last;
-
-    const cellId = isSetupCell ? SETUP_CELL_ID : CellId.create();
-    const undoCell = createCell({
-      id: cellId,
-      name,
-      code: serializedEditorState.doc,
-      edited: serializedEditorState.doc.trim().length > 0,
-      serializedEditorState,
-      config,
-    });
-
+    const toRestore = cellIds.map((id, i) => ({
+      id,
+      columnId: placements[i].columnId,
+      index: placements[i].index,
+    }));
     return {
       ...state,
-      cellIds: state.cellIds.insertId(cellId, column, index),
-      cellData: {
-        ...state.cellData,
-        [cellId]: undoCell,
-      },
-      cellRuntime: {
-        ...state.cellRuntime,
-        [cellId]: createCellRuntimeState(),
-      },
-      cellHandles: {
-        ...state.cellHandles,
-        [cellId]: createRef(),
-      },
-      history: state.history.slice(0, -1),
-      // Scroll to the newly created cell
-      scrollKey: cellId,
+      cellIds: state.cellIds.placeCells(toRestore),
+      scrollKey: cellIds[0] ?? null,
     };
   },
   clearSerializedEditorState: (state, action: { cellId: CellId }) => {
@@ -853,35 +643,6 @@ const {
       },
     });
   },
-  updateCellName: (state, action: { cellId: CellId; name: string }) => {
-    const { cellId, name } = action;
-    return updateCellData({
-      state,
-      cellId,
-      cellReducer: (cell) => {
-        return {
-          ...cell,
-          name: name,
-        };
-      },
-    });
-  },
-  updateCellConfig: (
-    state,
-    action: { cellId: CellId; config: Partial<CellConfig> },
-  ) => {
-    const { cellId, config } = action;
-    return updateCellData({
-      state,
-      cellId,
-      cellReducer: (cell) => {
-        return {
-          ...cell,
-          config: createCellConfig({ ...cell.config, ...config }),
-        };
-      },
-    });
-  },
   prepareForRun: (state, action: { cellId: CellId }) => {
     const newState = updateCellRuntimeState({
       state,
@@ -1020,97 +781,6 @@ const {
     }
     return { ...state, cellRuntime };
   },
-  /**
-   * Move focus to next cell
-   *
-   * Creates a new cell if the current cell is the last one in the array.
-   *
-   * If needed, scrolls newly created or focused cell into view.
-   *
-   * Replicates Shift+Enter functionality of Jupyter
-   */
-  moveToNextCell: (
-    state,
-    action: { cellId: CellId; before: boolean; noCreate?: boolean },
-  ) => {
-    const { cellId, before, noCreate = false } = action;
-
-    // Can't move focus of scratch cell
-    if (cellId === SCRATCH_CELL_ID) {
-      return state;
-    }
-
-    const column = state.cellIds.findWithId(cellId);
-    const index = column.indexOfOrThrow(cellId);
-    const nextCellIndex = before ? index - 1 : index + 1;
-
-    const isPastLastCell = nextCellIndex === column.length;
-    const isBeforeFirstCell = nextCellIndex === -1;
-
-    // Create a new cell at the end and set scrollKey to focus it
-    if (isPastLastCell && !noCreate) {
-      const newCellId = CellId.create();
-      return {
-        ...state,
-        cellIds: state.cellIds.insertId(newCellId, column.id, nextCellIndex),
-        cellData: {
-          ...state.cellData,
-          [newCellId]: createCell({ id: newCellId }),
-        },
-        cellRuntime: {
-          ...state.cellRuntime,
-          [newCellId]: createCellRuntimeState(),
-        },
-        cellHandles: {
-          ...state.cellHandles,
-          [newCellId]: createRef(),
-        },
-        scrollKey: newCellId,
-      };
-    }
-
-    if (isBeforeFirstCell && !noCreate) {
-      const newCellId = CellId.create();
-      return {
-        ...state,
-        cellIds: state.cellIds.insertId(newCellId, column.id, 0),
-        cellData: {
-          ...state.cellData,
-          [newCellId]: createCell({ id: newCellId }),
-        },
-        cellRuntime: {
-          ...state.cellRuntime,
-          [newCellId]: createCellRuntimeState(),
-        },
-        cellHandles: {
-          ...state.cellHandles,
-          [newCellId]: createRef(),
-        },
-        scrollKey: newCellId,
-      };
-    }
-
-    if ((isPastLastCell || isBeforeFirstCell) && noCreate) {
-      // Trying to move past the last cell or before the first cell
-      // with noCreate=true - do nothing
-      return state;
-    }
-
-    if (nextCellIndex >= 0 && nextCellIndex < column.length) {
-      // Move to valid cell within bounds
-      const nextCellId = column.atOrThrow(nextCellIndex);
-      // Just focus, no state change
-      focusAndScrollCellIntoView({
-        cellId: nextCellId,
-        cell: state.cellHandles[nextCellId],
-        isCodeHidden: isCellCodeHidden(state, nextCellId),
-        codeFocus: before ? "bottom" : "top",
-        variableName: undefined,
-      });
-    }
-
-    return state;
-  },
   markTouched: (state, action: { cellId: CellId }) => {
     const { cellId } = action;
 
@@ -1431,44 +1101,6 @@ const {
       cellRuntime: newCellRuntime,
     };
   },
-  addSetupCellIfDoesntExist: (state, action: { code?: string }) => {
-    const { code } = action;
-
-    if (state.cellIds.setupCellExists()) {
-      // Just focus on the existing setup cell
-      return {
-        ...state,
-        scrollKey: SETUP_CELL_ID,
-      };
-    }
-
-    return {
-      ...state,
-      cellIds: state.cellIds.insertId(
-        SETUP_CELL_ID,
-        state.cellIds.atOrThrow(0).id,
-        0,
-      ),
-      cellData: {
-        ...state.cellData,
-        [SETUP_CELL_ID]: createCell({
-          id: SETUP_CELL_ID,
-          name: SETUP_CELL_ID,
-          code,
-          edited: Boolean(code),
-        }),
-      },
-      cellRuntime: {
-        ...state.cellRuntime,
-        [SETUP_CELL_ID]: createCellRuntimeState(),
-      },
-      cellHandles: {
-        ...state.cellHandles,
-        [SETUP_CELL_ID]: createRef(),
-      },
-      scrollKey: SETUP_CELL_ID,
-    };
-  },
 });
 
 // We apply the middleware here (rather than inline in createReducerAndAtoms)
@@ -1876,20 +1508,13 @@ export function createTracebackInfoAtom(
   });
 }
 
-/**
- * Use this hook to dispatch cell actions. This hook will not cause a re-render
- * when cells change.
- */
-export function useCellActions(
-  options: { skipMiddleware?: boolean } = {},
-): CellActions {
-  return useActions(options);
+/** The reducer's actions. Components use `useCellActions` from `./actions`. */
+export function useReducerActions(): ReducerCellActions {
+  return useActions();
 }
 
-/**
- * Map of cell actions
- */
-export type CellActions = ReturnType<typeof createActions>;
+/** The actions the reducer still handles. */
+export type ReducerCellActions = ReturnType<typeof createActions>;
 
 /**
  * This is exported for testing purposes only.