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

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

@@ -21,7 +21,7 @@ const { stagedGenerationAtom, createActions, reducer, initialState } =
   visibleForTesting;
 
 // Mock the dependencies
-const mockCreateNewCell = vi.fn();
+const mockCreateCell = vi.fn();
 const mockUpdateCellCode = vi.fn();
 const mockDeleteCellCallback = vi.fn();
 
@@ -34,14 +34,16 @@ const mockCellHandle = {
   },
 };
 
+vi.mock("../../cells/actions", () => ({
+  useCellActions: () => ({
+    createCell: mockCreateCell,
+    updateCellCode: mockUpdateCellCode,
+  }),
+}));
 vi.mock("../../cells/cells", async () => {
   const { atom } = await import("jotai");
   return {
     notebookAtom: atom({ cellData: {}, cellIds: { inOrderIds: [] } }),
-    useCellActions: () => ({
-      createNewCell: mockCreateNewCell,
-      updateCellCode: mockUpdateCellCode,
-    }),
     cellHandleAtom: vi.fn(() => ({
       read: vi.fn(() => mockCellHandle),
     })),
@@ -261,11 +263,9 @@ describe("staged-cells", () => {
       const returnedCellId = result.current.createStagedCell(testCode);
 
       expect(returnedCellId).toBe(mockCellId);
-      expect(mockCreateNewCell).toHaveBeenCalledWith({
-        cellId: "__end__",
+      expect(mockCreateCell).toHaveBeenCalledWith({
         code: testCode,
-        before: false,
-        newCellId: mockCellId,
+        id: mockCellId,
       });
     });
 
@@ -330,7 +330,7 @@ describe("staged-cells", () => {
 
       // Verify it was created and added
       expect(createdCellId).toBe(mockCellId);
-      expect(mockCreateNewCell).toHaveBeenCalled();
+      expect(mockCreateCell).toHaveBeenCalled();
 
       let state = store.get(stagedAICellsAtom);
       expect(state.has(mockCellId)).toBe(true);
@@ -401,7 +401,7 @@ describe("staged cell generation", () => {
     result.current.beginStagedCellGeneration();
 
     // No cell or cell update should have been called
-    expect(mockCreateNewCell).not.toHaveBeenCalled();
+    expect(mockCreateCell).not.toHaveBeenCalled();
     expect(prepareCellCodeUpdate).not.toHaveBeenCalled();
     expect(store.get(stagedGenerationInProgressAtom)).toBe(true);
   });
@@ -610,11 +610,9 @@ describe("staged cell generation", () => {
       },
     });
 
-    expect(mockCreateNewCell).toHaveBeenCalledWith({
-      cellId: "__end__",
+    expect(mockCreateCell).toHaveBeenCalledWith({
       code: "print('```')",
-      before: false,
-      newCellId: "mock-cell-id",
+      id: "mock-cell-id",
     });
   });
 
@@ -636,16 +634,16 @@ describe("staged cell generation", () => {
       },
     });
 
-    expect(mockCreateNewCell).toHaveBeenCalledWith(
+    expect(mockCreateCell).toHaveBeenCalledWith(
       expect.objectContaining({
         code: "value = 1",
-        newCellId: "first-cell",
+        id: "first-cell",
       }),
     );
-    expect(mockCreateNewCell).toHaveBeenCalledWith(
+    expect(mockCreateCell).toHaveBeenCalledWith(
       expect.objectContaining({
         code: "value",
-        newCellId: "second-cell",
+        id: "second-cell",
       }),
     );
   });
@@ -683,7 +681,7 @@ describe("staged cell generation", () => {
       },
     });
 
-    expect(mockCreateNewCell).toHaveBeenCalledTimes(2);
+    expect(mockCreateCell).toHaveBeenCalledTimes(2);
     expect(store.get(stagedAICellsAtom).has(cellId("first-cell"))).toBe(true);
     expect(store.get(stagedAICellsAtom).has(cellId("second-cell"))).toBe(true);
     expect(prepareCellCodeUpdate).toHaveBeenCalledTimes(2);

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

@@ -14,11 +14,8 @@ import {
   notebookCellsCompletionSchema,
 } from "./completion-output";
 import { maybeAddMarimoImport } from "../cells/add-missing-import";
-import {
-  type CreateNewCellAction,
-  getCellEditorView,
-  useCellActions,
-} from "../cells/cells";
+import { getCellEditorView } from "../cells/cells";
+import { type CreateCell, useCellActions } from "../cells/actions";
 import { prepareCellCodeUpdate } from "../codemirror/language/cell-code";
 import type { JotaiStore } from "../state/jotai";
 import type { EditType } from "./tools/edit-notebook-tool";
@@ -98,18 +95,13 @@ interface UpdateStagedCellAction {
 /** Manage staged cells without assuming which AI workflow owns them. */
 export function useStagedCells(store: JotaiStore) {
   const { addStagedCell, removeStagedCell } = useStagedAICellsActions();
-  const { createNewCell, updateCellCode } = useCellActions();
+  const { updateCellCode, createCell } = useCellActions();
   const deleteCellCallback = useDeleteCellCallback();
 
   const createStagedCell = (code: string): CellId => {
     const newCellId = CellId.create();
     addStagedCell({ cellId: newCellId, edit: { type: "add_cell" } });
-    createNewCell({
-      cellId: "__end__",
-      code,
-      before: false,
-      newCellId: newCellId,
-    });
+    createCell({ code, id: newCellId });
     return newCellId;
   };
 
@@ -152,7 +144,7 @@ export function useStagedCells(store: JotaiStore) {
 /** Owns the lifecycle of cells created by one structured completion flow. */
 export function useStagedCellGeneration(store: JotaiStore) {
   const stagedCellActions = useStagedCells(store);
-  const { createNewCell } = useCellActions();
+  const { createCell } = useCellActions();
   const ownedGeneration = useRef<OwnedStagedGeneration | null>(null);
   // Ownership lives in a ref, so subscribe to its backing atoms to rerender
   // consumers when completion or per-cell acceptance changes.
@@ -180,7 +172,7 @@ export function useStagedCellGeneration(store: JotaiStore) {
       status: "in_progress",
       reconciler: new StagedCellReconciler({
         ...stagedCellActions,
-        createNewCell,
+        createCell,
       }),
     };
     store.set(stagedGenerationAtom, {
@@ -302,7 +294,7 @@ interface StagedCellReconcilerActions {
   deleteStagedCell: (cellId: CellId) => void;
   removeStagedCell: (cellId: CellId) => void;
   addStagedCell: (payload: { cellId: CellId; edit: Edit }) => void;
-  createNewCell: (opts: CreateNewCellAction) => void;
+  createCell: CreateCell;
 }
 
 class StagedCellReconciler {
@@ -391,7 +383,7 @@ class StagedCellReconciler {
 
     const cellId = maybeAddMarimoImport({
       autoInstantiate: false,
-      createNewCell: this.actions.createNewCell,
+      createCell: this.actions.createCell,
       fromCellId: this.createdCells[0]?.cellId,
       before: true,
     });

package/src/core/ai/tools/__tests__/edit-notebook-tool.test.ts

@@ -64,7 +64,7 @@ describe("EditNotebookTool", () => {
   let cellId3: CellId;
   let toolContext: {
     addStagedCell: ReturnType<typeof vi.fn>;
-    createNewCell: ReturnType<typeof vi.fn>;
+    createCell: ReturnType<typeof vi.fn>;
     prepareForRun: ReturnType<typeof vi.fn>;
     sendRun: ReturnType<typeof vi.fn>;
     store: ReturnType<typeof getDefaultStore>;
@@ -81,7 +81,7 @@ describe("EditNotebookTool", () => {
         const current = store.get(stagedAICellsAtom);
         store.set(stagedAICellsAtom, new Map(current).set(cellId, edit));
       }),
-      createNewCell: vi.fn(),
+      createCell: vi.fn(),
       prepareForRun: vi.fn(),
       sendRun: vi.fn().mockResolvedValue(null),
       store,

package/src/core/ai/tools/__tests__/run-cells-tool.test.ts

@@ -29,7 +29,7 @@ describe("RunStaleCellsTool", () => {
   let cellId3: CellId;
   let toolContext: {
     addStagedCell: ReturnType<typeof vi.fn>;
-    createNewCell: ReturnType<typeof vi.fn>;
+    createCell: ReturnType<typeof vi.fn>;
     prepareForRun: ReturnType<typeof vi.fn>;
     sendRun: ReturnType<typeof vi.fn>;
     store: ReturnType<typeof getDefaultStore>;
@@ -39,7 +39,7 @@ describe("RunStaleCellsTool", () => {
     store = getDefaultStore();
     toolContext = {
       addStagedCell: vi.fn(),
-      createNewCell: vi.fn(),
+      createCell: vi.fn(),
       prepareForRun: vi.fn(),
       sendRun: vi.fn().mockResolvedValue(null),
       store,

package/src/core/ai/tools/base.ts

@@ -1,7 +1,7 @@
 /* Copyright 2026 Marimo. All rights reserved. */
 
 import { z } from "zod";
-import type { CreateNewCellAction } from "@/core/cells/cells";
+import type { CreateCell } from "@/core/cells/actions";
 import type { CellId } from "@/core/cells/ids";
 import type { ExecuteCellsRequest } from "@/core/network/types";
 import type { JotaiStore } from "@/core/state/jotai";
@@ -113,7 +113,7 @@ export interface ToolDescription {
 export interface ToolNotebookContext {
   store: JotaiStore;
   addStagedCell: (payload: { cellId: CellId; edit: Edit }) => void;
-  createNewCell: (payload: CreateNewCellAction) => void;
+  createCell: CreateCell;
   prepareForRun: (payload: { cellId: CellId }) => void;
   sendRun: (request: ExecuteCellsRequest) => Promise<null>;
 }

package/src/core/ai/tools/edit-notebook-tool.ts

@@ -3,11 +3,8 @@
 import type { EditorView } from "@codemirror/view";
 import { z } from "zod";
 import { scrollAndHighlightCell } from "@/components/editor/links/cell-link";
-import {
-  type CellPosition as NotebookCellPosition,
-  type NotebookState,
-  notebookAtom,
-} from "@/core/cells/cells";
+import { type NotebookState, notebookAtom } from "@/core/cells/cells";
+import type { CreateCellOptions } from "@/core/cells/actions";
 import { CellId } from "@/core/cells/ids";
 import { prepareCellCodeUpdate } from "@/core/codemirror/language/cell-code";
 import type { CellColumnId } from "@/utils/id-tree";
@@ -86,7 +83,7 @@ export class EditNotebookTool implements AiTool<
     { edit }: EditNotebookInput,
     toolContext: ToolNotebookContext,
   ): Promise<ToolOutputBase> => {
-    const { addStagedCell, createNewCell, store } = toolContext;
+    const { addStagedCell, createCell, store } = toolContext;
 
     switch (edit.type) {
       case "update_cell": {
@@ -155,8 +152,7 @@ export class EditNotebookTool implements AiTool<
       case "add_cell": {
         const { position, code } = edit;
 
-        let notebookPosition: NotebookCellPosition = "__end__";
-        let before = false;
+        let placement: CreateCellOptions = {};
         const newCellId = CellId.create();
         const notebook = store.get(notebookAtom);
 
@@ -186,11 +182,12 @@ export class EditNotebookTool implements AiTool<
             }
 
             this.validateCellIdExists(cellId, notebook);
-            notebookPosition = cellId;
-            before = position.before;
+            placement = position.before
+              ? { before: cellId }
+              : { after: cellId };
             break;
           }
-          case "column_end": {
+          case "column_end":
             if (!position.columnIndex) {
               throw new ToolExecutionError(
                 "Column index is required for add_cell with column_end position",
@@ -203,21 +200,16 @@ export class EditNotebookTool implements AiTool<
               );
             }
 
-            const columnId = this.getColumnId(position.columnIndex, notebook);
-            notebookPosition = { type: "__end__", columnId };
+            placement = {
+              endOf: this.getColumnId(position.columnIndex, notebook),
+            };
             break;
-          }
+
           case "notebook_end":
-            // Use default: notebookPosition = "__end__"
             break;
         }
 
-        createNewCell({
-          cellId: notebookPosition,
-          before,
-          code,
-          newCellId,
-        });
+        createCell({ ...placement, code, id: newCellId });
 
         addStagedCell({
           cellId: newCellId,