@marimo-team/islands 0.25.3-dev209 → 0.25.3-dev210

package/src/core/ai/staged-cells.ts

@@ -95,7 +95,7 @@ interface UpdateStagedCellAction {
 /** Manage staged cells without assuming which AI workflow owns them. */
 export function useStagedCells(store: JotaiStore) {
   const { addStagedCell, removeStagedCell } = useStagedAICellsActions();
-  const { updateCellCode, createCell } = useCellActions();
+  const { setCellCode, createCell } = useCellActions();
   const deleteCellCallback = useDeleteCellCallback();
 
   const createStagedCell = (code: string): CellId => {
@@ -114,6 +114,8 @@ export function useStagedCells(store: JotaiStore) {
       return;
     }
 
+    // The AI may change the cell's language, and only the editor infers
+    // one. Without a mounted editor the code arrives through the document.
     const editorView = getCellEditorView(cellId);
     if (editorView) {
       const update = prepareCellCodeUpdate(editorView.state, code, {
@@ -123,7 +125,7 @@ export function useStagedCells(store: JotaiStore) {
         editorView.dispatch(update);
       }
     } else {
-      updateCellCode({ cellId, code, formattingChange: false });
+      setCellCode({ cellId, code });
     }
   };
 

package/src/core/cells/__tests__/cells.test.ts

@@ -178,6 +178,38 @@ describe("cell reducer", () => {
     return id;
   }
 
+  /** Edit a cell's code as the document delivers it. */
+  function editCode(cellId: CellId, code: string) {
+    const cells = state.cellIds.map((id) => {
+      const cell = state.cellData[id];
+      return {
+        id,
+        code: id === cellId ? code : cell.code,
+        name: cell.name,
+        config: {
+          column: cell.config.column ?? null,
+          disabled: cell.config.disabled ?? false,
+          hide_code: cell.config.hide_code ?? false,
+          expand_output: cell.config.expand_output ?? false,
+        },
+      };
+    });
+    actions.reconcileDocument({
+      order: state.cellIds,
+      cells: new Map(cells.map((cell) => [cell.id, cell])),
+    });
+  }
+
+  /** Format a cell's code, as the editor reports it. */
+  function formatCode(cellId: CellId, code: string) {
+    actions.codeFormatted({
+      cellId,
+      before: state.cellData[cellId].code,
+      after: code,
+    });
+    editCode(cellId, code);
+  }
+
   let i = 0;
   const originalCreate = CellId.create.bind(CellId);
 
@@ -200,11 +232,7 @@ describe("cell reducer", () => {
   });
 
   it("can update a cell", () => {
-    actions.updateCellCode({
-      cellId: firstCellId,
-      code: "import numpy as np",
-      formattingChange: false,
-    });
+    editCode(firstCellId, "import numpy as np");
     expect(formatCells(state)).toMatchInlineSnapshot(`
       "
       [0] 'import numpy as np'
@@ -243,11 +271,7 @@ describe("cell reducer", () => {
       stale_inputs: null,
       timestamp: 1,
     });
-    actions.updateCellCode({
-      cellId: firstCellId,
-      code: "print('new code')",
-      formattingChange: false,
-    });
+    editCode(firstCellId, "print('new code')");
     actions.prepareForRun({ cellId: firstCellId });
     actions.handleCellMessage({
       cell_id: firstCellId,
@@ -275,11 +299,7 @@ describe("cell reducer", () => {
     expect(cell).toMatchSnapshot("initial state");
 
     // Update code
-    actions.updateCellCode({
-      cellId: firstCellId,
-      code: "import marimo as mo",
-      formattingChange: false,
-    });
+    editCode(firstCellId, "import marimo as mo");
     cell = cells[0];
     expect(cell.status).toBe("idle");
     expect(cell.lastCodeRun).toBe(null);
@@ -384,32 +404,20 @@ describe("cell reducer", () => {
     // EDITING BACK AND FORTH
     /////////////////
     // Update code again
-    actions.updateCellCode({
-      cellId: firstCellId,
-      code: "import marimo as mo\nimport numpy",
-      formattingChange: false,
-    });
+    editCode(firstCellId, "import marimo as mo\nimport numpy");
     cell = cells[0];
     expect(cell.status).toBe("idle");
     expect(cell.edited).toBe(true);
     expect(cell).toMatchSnapshot("code edited again");
 
     // Update code should be unedited
-    actions.updateCellCode({
-      cellId: firstCellId,
-      code: "import marimo as mo",
-      formattingChange: false,
-    });
+    editCode(firstCellId, "import marimo as mo");
     cell = cells[0];
     expect(cell.edited).toBe(false);
     expect(cell).toMatchSnapshot("code reverted");
 
     // Update code should be edited again
-    actions.updateCellCode({
-      cellId: firstCellId,
-      code: "import marimo as mo\nimport numpy",
-      formattingChange: false,
-    });
+    editCode(firstCellId, "import marimo as mo\nimport numpy");
     cell = cells[0];
     expect(cell.status).toBe("idle");
     expect(cell.lastCodeRun).toBe("import marimo as mo");
@@ -533,11 +541,7 @@ describe("cell reducer", () => {
     expect(cell).toMatchSnapshot("initial state");
 
     // Update code
-    actions.updateCellCode({
-      cellId: firstCellId,
-      code: "import marimo as mo",
-      formattingChange: false,
-    });
+    editCode(firstCellId, "import marimo as mo");
     cell = cells[0];
     expect(cell.status).toBe("idle");
     expect(cell.lastCodeRun).toBe(null);
@@ -589,11 +593,7 @@ describe("cell reducer", () => {
 
   it("can run a stale cell", () => {
     // Update code of first cell
-    actions.updateCellCode({
-      cellId: firstCellId,
-      code: "import marimo as mo",
-      formattingChange: false,
-    });
+    editCode(firstCellId, "import marimo as mo");
     // Add cell
     seedCell({
       cellId: firstCellId,
@@ -601,11 +601,7 @@ describe("cell reducer", () => {
     });
     const secondCell = cells[1].id;
     // Update code
-    actions.updateCellCode({
-      code: "mo.slider()",
-      cellId: secondCell,
-      formattingChange: false,
-    });
+    editCode(secondCell, "mo.slider()");
 
     // Prepare for run
     actions.prepareForRun({
@@ -649,11 +645,7 @@ describe("cell reducer", () => {
 
   it("can format code and update cell", () => {
     const firstCellId = cells[0].id;
-    actions.updateCellCode({
-      cellId: firstCellId,
-      code: "import marimo as    mo",
-      formattingChange: false,
-    });
+    editCode(firstCellId, "import marimo as    mo");
     let cell = cells[0];
     expect(cell.status).toBe("idle");
     expect(cell.lastCodeRun).toBe(null);
@@ -689,11 +681,7 @@ describe("cell reducer", () => {
     expect(cell.lastCodeRun).toBe("import marimo as    mo");
 
     // Format code
-    actions.updateCellCode({
-      cellId: firstCellId,
-      code: "import marimo as mo",
-      formattingChange: true,
-    });
+    formatCode(firstCellId, "import marimo as mo");
     cell = cells[0];
     expect(cell.status).toBe("idle");
     expect(cell.lastCodeRun).toBe("import marimo as mo");
@@ -702,11 +690,7 @@ describe("cell reducer", () => {
 
   it("can run a stopped cell", () => {
     // Update code of first cell
-    actions.updateCellCode({
-      cellId: firstCellId,
-      code: "mo.md('This has an ancestor that was stopped')",
-      formattingChange: false,
-    });
+    editCode(firstCellId, "mo.md('This has an ancestor that was stopped')");
 
     // Prepare for run
     actions.prepareForRun({

package/src/core/cells/__tests__/document.test.ts

@@ -1,32 +1,17 @@
 /* Copyright 2026 Marimo. All rights reserved. */
 
-import { python } from "@codemirror/lang-python";
-import { EditorState } from "@codemirror/state";
-import { EditorView } from "@codemirror/view";
 import { atom } from "jotai";
 import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
 import { cellId } from "@/__tests__/branded";
 import type {
-  DocumentEdit,
-  DocumentEditor,
   DocumentModel,
   DocumentSyncStatus,
   EditableCell,
   EditableDocument,
 } from "@/core/notebook/document";
-import {
-  adaptiveLanguageConfiguration,
-  languageAdapterState,
-} from "@/core/codemirror/language/extension";
-import { OverridingHotkeyProvider } from "@/core/hotkeys/hotkeys";
 import { store } from "@/core/state/jotai";
 import { activeDocumentModelAtom } from "@/core/session/active-notebook";
-import {
-  createNotebookActions,
-  exportedForTesting,
-  type NotebookState,
-  notebookReducer,
-} from "../cells";
+import { exportedForTesting, type NotebookState } from "../cells";
 import { applyDocument, flushDocument, followDocument } from "../document";
 import { type CellId, SCRATCH_CELL_ID } from "../ids";
 
@@ -58,29 +43,6 @@ function documentOf(cells: EditableCell[]): EditableDocument {
   };
 }
 
-/** A mounted cell editor, as the notebook's cells have. */
-function createEditor(content: string) {
-  const state = EditorState.create({
-    doc: content,
-    extensions: [
-      python(),
-      adaptiveLanguageConfiguration({
-        cellId: a,
-        completionConfig: {
-          activate_on_typing: true,
-          signature_hint_on_typing: false,
-          copilot: false,
-          codeium_api_key: null,
-        },
-        hotkeys: new OverridingHotkeyProvider({}),
-        placeholderType: "marimo-import",
-        lspConfig: {},
-      }),
-    ],
-  });
-  return new EditorView({ state, parent: globalThis.document.body });
-}
-
 /** A notebook that follows `cells`, as the session builds one. */
 function notebookOf(cells: EditableCell[]): NotebookState {
   return applyDocument(initialNotebookState(), documentOf(cells));
@@ -110,38 +72,6 @@ describe("applyDocument", () => {
     expect(applyDocument(state, documentOf(cells))).toBe(state);
   });
 
-  it("brings a code change into the cell's editor", () => {
-    const state = notebookOf([cell(a), cell(b)]);
-    const view = createEditor("b = 1");
-    state.cellHandles[b].current = { editorView: view, editorViewOrNull: view };
-    state.cellData[b] = { ...state.cellData[b], lastCodeRun: "b = 1" };
-
-    const next = applyDocument(
-      state,
-      documentOf([cell(a), cell(b, { code: "b = 2" })]),
-    );
-
-    expect(next.cellData[b]).toMatchObject({ code: "b = 2", edited: true });
-    expect(next.cellData[a]).toBe(state.cellData[a]);
-    expect(view.state.doc.toString()).toBe("b = 2");
-  });
-
-  it("leaves an editor that already shows the document's code alone", () => {
-    const state = notebookOf([cell(a)]);
-    // The editor runs ahead of the notebook while this client types.
-    const view = createEditor('mo.md("hi")');
-    state.cellHandles[a].current = { editorView: view, editorViewOrNull: view };
-
-    const next = applyDocument(
-      state,
-      documentOf([cell(a, { code: 'mo.md("hi")' })]),
-    );
-
-    expect(next.cellData[a].code).toBe('mo.md("hi")');
-    expect(view.state.doc.toString()).toBe('mo.md("hi")');
-    expect(view.state.field(languageAdapterState).type).toBe("python");
-  });
-
   it("keeps the code this client has when the host withholds source", () => {
     const state = notebookOf([cell(a)]);
 
@@ -182,74 +112,6 @@ describe("applyDocument", () => {
   });
 });
 
-describe("documentEditMiddleware", () => {
-  type Call = [string, ...unknown[]];
-  let calls: Call[];
-  let state: NotebookState;
-  const editor: DocumentEditor = {
-    setCode: (id, code) => calls.push(["setCode", id, code]),
-    setName: (id, name) => calls.push(["setName", id, name]),
-    updateConfig: (id, patch) => calls.push(["updateConfig", id, patch]),
-    createCell: (cell, placement) =>
-      calls.push(["createCell", cell, placement]),
-    deleteCell: (id) => calls.push(["deleteCell", id]),
-    setOrder: (order) => calls.push(["setOrder", order]),
-  };
-  const model: DocumentModel = {
-    documentAtom: atom(null),
-    editAtom: atom(null, (_get, _set, edit: DocumentEdit) => edit(editor)),
-    syncStatusAtom: atom<DocumentSyncStatus>({ status: "idle" }),
-    flush: vi.fn(),
-    retry: vi.fn(),
-    start: vi.fn(),
-  };
-  const actions = createNotebookActions((action) => {
-    state = notebookReducer(state, action);
-  });
-
-  beforeEach(() => {
-    calls = [];
-    state = notebookOf([cell(a), cell(b), cell(c)]);
-    store.set(activeDocumentModelAtom, model);
-  });
-
-  afterEach(() => {
-    store.set(activeDocumentModelAtom, null);
-  });
-
-  it("writes code", () => {
-    actions.updateCellCode({
-      cellId: a,
-      code: "a = 2",
-      formattingChange: false,
-    });
-
-    expect(calls).toEqual([["setCode", a, "a = 2"]]);
-  });
-
-  it("ignores the document and the kernel's run state", () => {
-    actions.prepareForRun({ cellId: a });
-    actions.setLastRun({
-      lastCodeRun: { [a]: "a = 0" },
-      lastExecutionTime: {},
-    });
-    actions.resetCellRuntimes();
-    actions.reconcileDocument(documentOf([cell(a)]));
-
-    expect(calls).toEqual([]);
-  });
-
-  it("leaves the scratchpad alone", () => {
-    actions.updateCellCode({
-      cellId: SCRATCH_CELL_ID,
-      code: "scratch",
-      formattingChange: false,
-    });
-
-    expect(calls).toEqual([]);
-  });
-});
-
 describe("followDocument", () => {
   it("reconciles the document now and on each change, until aborted", () => {
     const documentAtom = atom<EditableDocument | null>(null);

package/src/core/cells/__tests__/notebook-outline.test.ts

@@ -42,16 +42,18 @@ function setup() {
 
 describe("notebookOutline", () => {
   it("does not notify outline subscribers when cell code changes", () => {
-    const { store, actions } = setup();
+    const { store } = setup();
     const initial = store.get(notebookOutline);
     const listener = vi.fn();
     const unsubscribe = store.sub(notebookOutline, listener);
 
-    actions.updateCellCode({
-      cellId: first,
-      code: "x = 2",
-      formattingChange: false,
-    });
+    store.set(notebookAtom, (state) => ({
+      ...state,
+      cellData: {
+        ...state.cellData,
+        [first]: { ...state.cellData[first], code: "x = 2" },
+      },
+    }));
 
     expect(store.get(notebookOutline)).toBe(initial);
     expect(listener).not.toHaveBeenCalled();

package/src/core/cells/actions.ts

@@ -159,6 +159,20 @@ export const createCellAtom = atom(
   },
 );
 
+/**
+ * Write a cell's code, or nothing for a cell the document no longer has.
+ * A mounted editor follows the document.
+ */
+export const setCellCodeAtom = atom(
+  null,
+  (get, set, { cellId, code }: { cellId: CellId; code: string }): void => {
+    if (!documentOf(get).cells.has(cellId)) {
+      return;
+    }
+    set(modelOf(get).editAtom, (tx) => tx.setCode(cellId, code));
+  },
+);
+
 /**
  * Split a cell at its cursor. The text before stays, the text after
  * becomes a new cell below, which is focused. Returns the new cell's id,
@@ -398,6 +412,27 @@ export const addSetupCellAtom = atom(
   },
 );
 
+/**
+ * Create the setup cell with `code`, or append `code` to the one that exists.
+ */
+export const addOrAppendSetupCellAtom = atom(
+  null,
+  (get, set, { code }: { code: string }): void => {
+    const setup = documentOf(get).cells.get(SETUP_CELL_ID);
+    if (setup === undefined) {
+      set(addSetupCellAtom, { code });
+      return;
+    }
+    // Withheld source keeps the code this client has.
+    const existing =
+      setup.code ?? get(notebookAtom).cellData[SETUP_CELL_ID]?.code ?? "";
+    set(setCellCodeAtom, {
+      cellId: SETUP_CELL_ID,
+      code: existing.trim() ? `${existing.trimEnd()}\n\n${code}` : code,
+    });
+  },
+);
+
 /**
  * Delete cells, keeping them in the history for undo. The notebook keeps at
  * least one cell. Returns false when nothing was deleted.
@@ -574,6 +609,7 @@ export function useCellActions() {
   const createCell = useSetAtom(createCellAtom);
   const moveToNextCell = useSetAtom(moveToNextCellAtom);
   const addSetupCell = useSetAtom(addSetupCellAtom);
+  const addOrAppendSetupCell = useSetAtom(addOrAppendSetupCellAtom);
   const deleteCells = useSetAtom(deleteCellsAtom);
   const undoDeleteCell = useSetAtom(undoDeleteCellAtom);
   const focusCell = useSetAtom(focusCellAtom);
@@ -597,11 +633,13 @@ export function useCellActions() {
   const mergeAllColumns = useSetAtom(mergeAllColumnsAtom);
   const compactColumns = useSetAtom(compactColumnsAtom);
   const moveColumn = useSetAtom(moveColumnAtom);
+  const setCellCode = useSetAtom(setCellCodeAtom);
   const splitCell = useSetAtom(splitCellAtom);
   const undoSplitCell = useSetAtom(undoSplitCellAtom);
   return useMemo(
     () => ({
       ...operations,
+      setCellCode,
       splitCell,
       undoSplitCell,
       addColumn,
@@ -620,6 +658,7 @@ export function useCellActions() {
       createCell,
       moveToNextCell,
       addSetupCell,
+      addOrAppendSetupCell,
       deleteCells,
       undoDeleteCell,
       focusCell,

package/src/core/cells/cells.ts

@@ -22,7 +22,7 @@ import type { CellConfig } from "../network/types";
 import { createDeepEqualAtom, store } from "../state/jotai";
 import { isWasm } from "../wasm/utils";
 import { prepareCellForExecution, transitionCell } from "./cell";
-import { applyDocument, documentEditMiddleware } from "./document";
+import { applyDocument } from "./document";
 import { isCodeHidden, shownHiddenCodeAtom } from "./editor";
 import { type CellId, SCRATCH_CELL_ID } from "./ids";
 import { type CellLog, getCellLogsForMessage } from "./logs";
@@ -154,50 +154,35 @@ export function initialNotebookState(): NotebookState {
  */
 const {
   reducer,
-  addMiddleware,
   createActions,
   useActions,
   valueAtom: notebookAtom,
 } = createReducerAndAtoms(initialNotebookState, {
   /**
-   * Put moved cells back where they were.
-   *
-   * TODO(MO-8028): Move to `undoDeleteCellAtom` with the order family.
+   * Keep a cell as run, or as edited, as it was before formatting rewrote
+   * its code from `before` to `after`. The code itself arrives through the
+   * document.
    */
-  updateCellCode: (
+  codeFormatted: (
     state,
-    action: {
-      cellId: CellId;
-      code: string;
-      /**
-       * Whether or not the update is a formatting change,
-       * if so, the 'edited' state will be handled differently.
-       */
-      formattingChange: boolean;
-    },
+    action: { cellId: CellId; before: string; after: string },
   ) => {
-    const { cellId, code, formattingChange } = action;
+    const { cellId, before, after } = action;
     if (!state.cellData[cellId]) {
       return state;
     }
-
     return updateCellData({
       state,
       cellId,
       cellReducer: (cell) => {
-        // Formatting-only change means we can re-use the last code run
-        // if it was not previously edited. And we don't change the edited state.
-        return formattingChange
-          ? {
-              ...cell,
-              code: code,
-              lastCodeRun: cell.edited ? cell.lastCodeRun : code,
-            }
-          : {
-              ...cell,
-              code: code,
-              edited: code.trim() !== cell.lastCodeRun,
-            };
+        const lastCodeRun = isEdited(cell.lastCodeRun, before)
+          ? cell.lastCodeRun
+          : after;
+        return {
+          ...cell,
+          lastCodeRun,
+          edited: isEdited(lastCodeRun, cell.code),
+        };
       },
     });
   },
@@ -426,11 +411,6 @@ const {
   },
 });
 
-// We apply the middleware here (rather than inline in createReducerAndAtoms)
-// so that the document edit middleware can import CellActions and
-// strictly type the dispatched actions without creating a circular dependency.
-addMiddleware(documentEditMiddleware);
-
 // Helper function to update a cell in the array
 function updateCellRuntimeState({
   state,