@marimo-team/islands 0.25.3-dev28 → 0.25.3-dev284

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

@@ -4,13 +4,18 @@ import { EditorState } from "@codemirror/state";
 import { EditorView } from "@codemirror/view";
 import { getDefaultStore } from "jotai";
 import { beforeEach, describe, expect, it, vi } from "vitest";
+import { mockDocumentModel } from "@/__mocks__/document";
 import { MockNotebook } from "@/__mocks__/notebook";
 import { cellId } from "@/__tests__/branded";
 import { notebookAtom } from "@/core/cells/cells";
 import type { CellId } from "@/core/cells/ids";
-import { updateEditorCodeFromPython } from "@/core/codemirror/language/utils";
 import { OverridingHotkeyProvider } from "@/core/hotkeys/hotkeys";
-import { MultiColumn } from "@/utils/id-tree";
+import {
+  cell,
+  CONFIG,
+  document as documentOf,
+  withCells,
+} from "@/core/notebook/document/__tests__/fixtures";
 import { cellConfigExtension } from "../../../codemirror/config/extension";
 import { adaptiveLanguageConfiguration } from "../../../codemirror/language/extension";
 import { stagedAICellsAtom } from "../../staged-cells";
@@ -22,11 +27,6 @@ vi.mock("@/components/editor/links/cell-link", () => ({
   scrollAndHighlightCell: vi.fn(),
 }));
 
-// Mock updateEditorCodeFromPython
-vi.mock("@/core/codemirror/language/utils", () => ({
-  updateEditorCodeFromPython: vi.fn(),
-}));
-
 function createMockEditorView(code: string): EditorView {
   return new EditorView({
     state: EditorState.create({
@@ -67,12 +67,9 @@ describe("EditNotebookTool", () => {
   let tool: EditNotebookTool;
   let cellId1: CellId;
   let cellId2: CellId;
-  let cellId3: CellId;
   let toolContext: {
     addStagedCell: ReturnType<typeof vi.fn>;
-    createNewCell: ReturnType<typeof vi.fn>;
-    prepareForRun: ReturnType<typeof vi.fn>;
-    sendRun: ReturnType<typeof vi.fn>;
+    createCell: ReturnType<typeof vi.fn>;
     store: ReturnType<typeof getDefaultStore>;
   };
 
@@ -87,16 +84,13 @@ describe("EditNotebookTool", () => {
         const current = store.get(stagedAICellsAtom);
         store.set(stagedAICellsAtom, new Map(current).set(cellId, edit));
       }),
-      createNewCell: vi.fn(),
-      prepareForRun: vi.fn(),
-      sendRun: vi.fn().mockResolvedValue(null),
+      createCell: vi.fn(),
       store,
     };
     tool = new EditNotebookTool();
 
     cellId1 = cellId("cell-1");
     cellId2 = cellId("cell-2");
-    cellId3 = cellId("cell-3");
 
     // Reset mocks
     vi.clearAllMocks();
@@ -106,7 +100,7 @@ describe("EditNotebookTool", () => {
     it("should have correct metadata", () => {
       expect(tool.name).toBe("edit_notebook_tool");
       expect(tool.description).toBeDefined();
-      expect(tool.description.baseDescription).toContain("editing operations");
+      expect(tool.description.baseDescription).toContain("Propose cell edits");
       expect(tool.schema).toBeDefined();
       expect(tool.outputSchema).toBeDefined();
       expect(tool.mode).toEqual(["agent"]);
@@ -140,10 +134,7 @@ describe("EditNotebookTool", () => {
       );
 
       expect(result.status).toBe("success");
-      expect(vi.mocked(updateEditorCodeFromPython)).toHaveBeenCalledWith(
-        editorView,
-        newCode,
-      );
+      expect(editorView.state.doc.toString()).toBe(newCode);
 
       // Check that cell was staged
       const stagedCells = store.get(stagedAICellsAtom);
@@ -166,10 +157,7 @@ describe("EditNotebookTool", () => {
       );
 
       expect(result2.status).toBe("success");
-      expect(vi.mocked(updateEditorCodeFromPython)).toHaveBeenCalledWith(
-        editorView,
-        "x = 3",
-      );
+      expect(editorView.state.doc.toString()).toBe("x = 3");
 
       const stagedCells2 = store.get(stagedAICellsAtom);
       expect(stagedCells2.has(cellId1)).toBe(true);
@@ -206,7 +194,7 @@ describe("EditNotebookTool", () => {
 
       // Should NOT stage the cell or update the editor
       expect(toolContext.addStagedCell).not.toHaveBeenCalled();
-      expect(vi.mocked(updateEditorCodeFromPython)).not.toHaveBeenCalled();
+      expect(editorView.state.doc.toString()).toBe(code);
 
       // No staged cells
       const stagedCells = store.get(stagedAICellsAtom);
@@ -459,27 +447,37 @@ describe("EditNotebookTool", () => {
           [cellId2]: { code: "x = 2" },
         },
       });
-      // Create multi-column layout
-      notebook.cellIds = MultiColumn.from([[cellId1], [cellId2]]);
       store.set(notebookAtom, notebook);
-
-      const newCode = "y = 2";
-      const result = await tool.handler(
-        {
-          edit: {
-            type: "add_cell",
-            position: { type: "column_end", columnIndex: 1 },
-            code: newCode,
-          },
-        },
-        toolContext as never,
+      // The layout reads columns from the document's cells.
+      const { stop } = mockDocumentModel(
+        store,
+        withCells(documentOf(), {
+          c1: null,
+          c2: null,
+          [cellId1]: cell(cellId1, "x = 1"),
+          [cellId2]: cell(cellId2, "x = 2", {
+            config: { ...CONFIG, column: 1 },
+          }),
+        }),
       );
 
-      expect(result.status).toBe("success");
+      try {
+        const result = await tool.handler(
+          {
+            edit: {
+              type: "add_cell",
+              position: { type: "column_end", columnIndex: 1 },
+              code: "y = 2",
+            },
+          },
+          toolContext as never,
+        );
 
-      // Check that a new cell was staged
-      const stagedCells = store.get(stagedAICellsAtom);
-      expect(stagedCells.size).toBe(1);
+        expect(result.status).toBe("success");
+        expect(store.get(stagedAICellsAtom).size).toBe(1);
+      } finally {
+        stop();
+      }
     });
 
     it("should throw error when cell ID doesn't exist for position", async () => {
@@ -714,52 +712,6 @@ describe("EditNotebookTool", () => {
     });
   });
 
-  describe("validation", () => {
-    it("should validate cell exists in multi-column notebook", async () => {
-      const notebook = MockNotebook.notebookState({
-        cellData: {
-          [cellId1]: { code: "x = 1" },
-          [cellId2]: { code: "x = 2" },
-          [cellId3]: { code: "x = 3" },
-        },
-      });
-      // Create multi-column layout
-      notebook.cellIds = MultiColumn.from([[cellId1, cellId2], [cellId3]]);
-      store.set(notebookAtom, notebook);
-
-      // Should not throw for cells in different columns
-      const editorView = createMockEditorView("x = 1");
-      notebook.cellHandles[cellId1] = { current: { editorView } } as never;
-      notebook.cellHandles[cellId3] = { current: { editorView } } as never;
-
-      await expect(
-        tool.handler(
-          {
-            edit: {
-              type: "update_cell",
-              position: { cellId: cellId1 },
-              code: "y = 1",
-            },
-          },
-          toolContext as never,
-        ),
-      ).resolves.toBeDefined();
-
-      await expect(
-        tool.handler(
-          {
-            edit: {
-              type: "update_cell",
-              position: { cellId: cellId3 },
-              code: "y = 3",
-            },
-          },
-          toolContext as never,
-        ),
-      ).resolves.toBeDefined();
-    });
-  });
-
   describe("return value", () => {
     it("should return success status with next steps", async () => {
       const notebook = MockNotebook.notebookState({

package/src/core/ai/tools/__tests__/registry.test.ts

@@ -10,6 +10,23 @@ import { TestFrontendTool } from "../sample-tool";
 const INVALID_KEYS_IN_SCHEMA_PARAMS = new Set(["anyOf", "oneOf"]);
 
 describe("FrontendToolRegistry", () => {
+  it("only offers cell proposals with Suggest edits access", () => {
+    expect(
+      Object.fromEntries(
+        COPILOT_MODES.map((mode) => [
+          mode,
+          FRONTEND_TOOL_REGISTRY.getToolSchemas(mode).map((tool) => tool.name),
+        ]),
+      ),
+    ).toEqual({
+      manual: [],
+      ask: [],
+      agent: ["edit_notebook_tool"],
+      code_mode: [],
+    });
+    expect(FRONTEND_TOOL_REGISTRY.has("run_stale_cells_tool")).toBe(false);
+  });
+
   it("registers tools via constructor and supports has()", () => {
     const registry = new FrontendToolRegistry([new TestFrontendTool()]);
     expect(registry.has("test_frontend_tool")).toBe(true);

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

@@ -1,9 +1,8 @@
 /* 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";
 import type { Edit } from "../staged-cells";
 import type { CopilotMode } from "./registry";
@@ -20,9 +19,11 @@ export type StatusValue = "success" | "error" | "warning";
  * @example
  * throw new ToolExecutionError(
  *   "Failed to fetch data",
- *   "FETCH_ERROR",
- *   true,
- *   "Check your network connection"
+ *   {
+ *     code: "FETCH_ERROR",
+ *     isRetryable: true,
+ *     suggestedFix: "Check your network connection",
+ *   }
  * );
  */
 export class ToolExecutionError extends Error {
@@ -33,10 +34,17 @@ export class ToolExecutionError extends Error {
 
   public constructor(
     message: string,
-    code = "TOOL_ERROR",
-    isRetryable = false,
-    suggestedFix?: string,
-    meta?: Record<string, unknown>,
+    {
+      code = "TOOL_ERROR",
+      isRetryable = false,
+      suggestedFix,
+      meta,
+    }: {
+      code?: string;
+      isRetryable?: boolean;
+      suggestedFix?: string;
+      meta?: Record<string, unknown>;
+    } = {},
   ) {
     super(message);
     this.name = "ToolExecutionError";
@@ -58,9 +66,6 @@ export class ToolExecutionError extends Error {
   }
 }
 
-// Can be any input
-export type EmptyToolInput = unknown;
-
 /**
  * Base interface for tool output responses.
  * Mirrors the SuccessResult dataclass from marimo/_ai/_tools/types.py
@@ -104,9 +109,7 @@ export interface ToolDescription {
 export interface ToolNotebookContext {
   store: JotaiStore;
   addStagedCell: (payload: { cellId: CellId; edit: Edit }) => void;
-  createNewCell: (payload: CreateNewCellAction) => void;
-  prepareForRun: (payload: { cellId: CellId }) => void;
-  sendRun: (request: ExecuteCellsRequest) => Promise<null>;
+  createCell: CreateCell;
 }
 
 export interface AiTool<TIn, TOut> {

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

@@ -3,14 +3,15 @@
 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 { updateEditorCodeFromPython } from "@/core/codemirror/language/utils";
-import type { CellColumnId } from "@/utils/id-tree";
+import { prepareCellCodeUpdate } from "@/core/codemirror/language/cell-code";
+import {
+  type CellColumnId,
+  type Layout,
+  layoutAtom,
+} from "@/core/cells/layout";
 import { stagedAICellsAtom } from "../staged-cells";
 import {
   type AiTool,
@@ -24,7 +25,7 @@ import type { CopilotMode } from "./registry";
 
 const description: ToolDescription = {
   baseDescription:
-    "Perform editing operations on the current notebook. You should prefer to create new cells unless you need to edit existing cells. Call this tool multiple times to perform multiple edits. Separate code into logical individual cells to take advantage of the notebook's reactive execution model.",
+    "Propose cell edits for the user to review in the current notebook. You should prefer to create new cells unless you need to edit existing cells. Call this tool multiple times to perform multiple edits. Separate code into logical individual cells to take advantage of the notebook's reactive execution model.",
   prerequisites: [
     "If you are updating existing cells, you need the cellIds. If they are not known, call the lightweight_cell_map_tool to find out.",
   ],
@@ -86,7 +87,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": {
@@ -94,19 +95,17 @@ export class EditNotebookTool implements AiTool<
         const cellId = position.cellId as CellId | undefined;
 
         if (!cellId) {
-          throw new ToolExecutionError(
-            "Cell ID is required for update_cell",
-            "CELL_ID_REQUIRED_FOR_UPDATE_CELL",
-            true,
-            "Provide a cell ID to update",
-          );
+          throw new ToolExecutionError("Cell ID is required for update_cell", {
+            code: "CELL_ID_REQUIRED_FOR_UPDATE_CELL",
+            isRetryable: true,
+            suggestedFix: "Provide a cell ID to update",
+          });
         } else if (!code) {
-          throw new ToolExecutionError(
-            "Code is required for update_cell",
-            "CODE_REQUIRED_FOR_UPDATE_CELL",
-            true,
-            "Provide the new code to update the cell",
-          );
+          throw new ToolExecutionError("Code is required for update_cell", {
+            code: "CODE_REQUIRED_FOR_UPDATE_CELL",
+            isRetryable: true,
+            suggestedFix: "Provide the new code to update the cell",
+          });
         }
 
         const notebook = store.get(notebookAtom);
@@ -119,7 +118,12 @@ export class EditNotebookTool implements AiTool<
 
         // If previous edit was from a new cell, just replace editor code with the new code
         if (existingStagedCell?.type === "add_cell") {
-          updateEditorCodeFromPython(editorView, code);
+          const update = prepareCellCodeUpdate(editorView.state, code, {
+            initialize: true,
+          });
+          if (update) {
+            editorView.dispatch(update);
+          }
           break;
         }
 
@@ -139,15 +143,20 @@ export class EditNotebookTool implements AiTool<
           edit: { type: "update_cell", previousCode: previousCode },
         });
 
-        updateEditorCodeFromPython(editorView, code);
+        const update = prepareCellCodeUpdate(editorView.state, code, {
+          initialize: true,
+        });
+
+        if (update) {
+          editorView.dispatch(update);
+        }
 
         break;
       }
       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);
 
@@ -157,51 +166,57 @@ export class EditNotebookTool implements AiTool<
             if (!cellId) {
               throw new ToolExecutionError(
                 "Cell ID is required for add_cell with relative position",
-                "TOOL_ERROR",
-                true,
-                "Provide a cell ID to add the cell before",
+                {
+                  code: "TOOL_ERROR",
+                  isRetryable: true,
+                  suggestedFix: "Provide a cell ID to add the cell before",
+                },
               );
             }
 
             if (position.before === undefined) {
               throw new ToolExecutionError(
                 "Before is required for add_cell with relative position",
-                "TOOL_ERROR",
-                true,
-                `Provide a boolean value for before, true to add before cell ${cellId}, false to add after the cell`,
+                {
+                  code: "TOOL_ERROR",
+                  isRetryable: true,
+                  suggestedFix: `Provide a boolean value for before, true to add before cell ${cellId}, false to add after the cell`,
+                },
               );
             }
 
             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",
-                "TOOL_ERROR",
-                true,
-                "Provide a column index to add the cell at the end of",
+                {
+                  code: "TOOL_ERROR",
+                  isRetryable: true,
+                  suggestedFix:
+                    "Provide a column index to add the cell at the end of",
+                },
               );
             }
 
-            const columnId = this.getColumnId(position.columnIndex, notebook);
-            notebookPosition = { type: "__end__", columnId };
+            placement = {
+              endOf: this.getColumnId(
+                position.columnIndex,
+                store.get(layoutAtom),
+              ),
+            };
             break;
-          }
+
           case "notebook_end":
-            // Use default: notebookPosition = "__end__"
             break;
         }
 
-        createNewCell({
-          cellId: notebookPosition,
-          before,
-          code,
-          newCellId,
-        });
+        createCell({ ...placement, code, id: newCellId });
 
         addStagedCell({
           cellId: newCellId,
@@ -218,12 +233,11 @@ export class EditNotebookTool implements AiTool<
         const cellId = position.cellId as CellId | undefined;
 
         if (!cellId) {
-          throw new ToolExecutionError(
-            "Cell ID is required for delete_cell",
-            "CELL_ID_REQUIRED_FOR_DELETE_CELL",
-            true,
-            "Provide a cell ID to delete",
-          );
+          throw new ToolExecutionError("Cell ID is required for delete_cell", {
+            code: "CELL_ID_REQUIRED_FOR_DELETE_CELL",
+            isRetryable: true,
+            suggestedFix: "Provide a cell ID to delete",
+          });
         }
 
         const notebook = store.get(notebookAtom);
@@ -247,37 +261,30 @@ export class EditNotebookTool implements AiTool<
       next_steps: [
         "If you need to perform more edits, call this tool again.",
         "You should use the lint notebook tool to check for errors and lint issues. Fix them by editing the notebook.",
-        "You should use the run stale cells tool to run the cells that have been edited or newly added. This allows you to see the output of the cells and fix any errors.",
       ],
     };
   };
 
   private validateCellIdExists(cellId: CellId, notebook: NotebookState) {
-    const cellIds = notebook.cellIds;
-    if (!cellIds.getColumns().some((column) => column.idSet.has(cellId))) {
-      throw new ToolExecutionError(
-        "Cell not found",
-        "CELL_NOT_FOUND",
-        false,
-        "Check which cells exist in the notebook",
-      );
+    if (!notebook.cellIds.includes(cellId)) {
+      throw new ToolExecutionError("Cell not found", {
+        code: "CELL_NOT_FOUND",
+        isRetryable: false,
+        suggestedFix: "Check which cells exist in the notebook",
+      });
     }
   }
 
-  private getColumnId(
-    columnIndex: number,
-    notebook: NotebookState,
-  ): CellColumnId {
-    const cellIds = notebook.cellIds;
-    const columns = cellIds.getColumns();
+  private getColumnId(columnIndex: number, layout: Layout): CellColumnId {
+    const { columns } = layout;
 
     if (columnIndex < 0 || columnIndex >= columns.length) {
-      throw new ToolExecutionError(
-        "Column index is out of range",
-        "COLUMN_INDEX_OUT_OF_RANGE",
-        true,
-        "Choose a column index between 0 and the number of columns in the notebook (0-based)",
-      );
+      throw new ToolExecutionError("Column index is out of range", {
+        code: "COLUMN_INDEX_OUT_OF_RANGE",
+        isRetryable: true,
+        suggestedFix:
+          "Choose a column index between 0 and the number of columns in the notebook (0-based)",
+      });
     }
     return columns[columnIndex].id;
   }
@@ -289,12 +296,11 @@ export class EditNotebookTool implements AiTool<
     const cellHandles = notebook.cellHandles;
     const cellHandle = cellHandles[cellId].current;
     if (!cellHandle?.editorView) {
-      throw new ToolExecutionError(
-        "Cell editor not found",
-        "CELL_EDITOR_NOT_FOUND",
-        false,
-        "Internal error, ask the user to report this error",
-      );
+      throw new ToolExecutionError("Cell editor not found", {
+        code: "CELL_EDITOR_NOT_FOUND",
+        isRetryable: false,
+        suggestedFix: "Internal error, ask the user to report this error",
+      });
     }
     return cellHandle.editorView;
   }

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

@@ -9,7 +9,6 @@ import {
   type ToolNotebookContext,
 } from "./base";
 import { EditNotebookTool } from "./edit-notebook-tool";
-import { RunStaleCellsTool } from "./run-cells-tool";
 import { formatToolDescription } from "./utils";
 
 export type AnyZodObject = ZodObject<z.ZodRawShape>;
@@ -73,9 +72,11 @@ export class FrontendToolRegistry {
         const strError = z.prettifyError(inputResponse.error);
         throw new ToolExecutionError(
           `Tool ${toolName} returned invalid input: ${strError}`,
-          "INVALID_ARGUMENTS",
-          true,
-          "Please check the arguments and try again.",
+          {
+            code: "INVALID_ARGUMENTS",
+            isRetryable: true,
+            suggestedFix: "Please check the arguments and try again.",
+          },
         );
       }
       const args = inputResponse.data;
@@ -110,12 +111,15 @@ export class FrontendToolRegistry {
       // Handle unknown/generic errors
       const genericError = new ToolExecutionError(
         error instanceof Error ? error.message : String(error),
-        "TOOL_ERROR",
-        false,
-        "Check the error message and try again with valid arguments.",
         {
-          args: rawArgs,
-          errorType: error?.constructor?.name ?? typeof error,
+          code: "TOOL_ERROR",
+          isRetryable: false,
+          suggestedFix:
+            "Check the error message and try again with valid arguments.",
+          meta: {
+            args: rawArgs,
+            errorType: error?.constructor?.name ?? typeof error,
+          },
         },
       );
       return {
@@ -143,5 +147,4 @@ export class FrontendToolRegistry {
 
 export const FRONTEND_TOOL_REGISTRY = new FrontendToolRegistry([
   new EditNotebookTool(),
-  new RunStaleCellsTool(),
 ]);

package/src/core/ai/tools/sample-tool.ts

@@ -50,13 +50,12 @@ export class TestFrontendTool implements AiTool<Input, Output> {
   public async handler({ name }: Input): Promise<Output> {
     // Example: Validate input and throw ToolExecutionError on invalid data
     if (!name.trim()) {
-      throw new ToolExecutionError(
-        "Name cannot be empty",
-        "INVALID_INPUT",
-        false,
-        "Please provide a non-empty name",
-        { field: "name", received: name },
-      );
+      throw new ToolExecutionError("Name cannot be empty", {
+        code: "INVALID_INPUT",
+        isRetryable: false,
+        suggestedFix: "Please provide a non-empty name",
+        meta: { field: "name", received: name },
+      });
     }
 
     return {