@marimo-team/islands 0.25.3-dev24 → 0.25.3-dev244

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 {

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

@@ -23,7 +23,6 @@ exports[`cell reducer > can initialize stdin > initial state 1`] = `
   "output": null,
   "runElapsedTimeMs": null,
   "runStartTimestamp": null,
-  "serializedEditorState": null,
   "staleInputs": false,
   "status": "idle",
   "stopped": false,
@@ -53,7 +52,6 @@ exports[`cell reducer > can initialize stdin > prepared for run 1`] = `
   "output": null,
   "runElapsedTimeMs": null,
   "runStartTimestamp": null,
-  "serializedEditorState": null,
   "staleInputs": false,
   "status": "queued",
   "stopped": false,
@@ -83,7 +81,6 @@ exports[`cell reducer > can initialize stdin > queued message received 1`] = `
   "output": null,
   "runElapsedTimeMs": null,
   "runStartTimestamp": null,
-  "serializedEditorState": null,
   "staleInputs": false,
   "status": "queued",
   "stopped": false,
@@ -113,7 +110,6 @@ exports[`cell reducer > can initialize stdin > running message received 1`] = `
   "output": null,
   "runElapsedTimeMs": null,
   "runStartTimestamp": 20,
-  "serializedEditorState": null,
   "staleInputs": false,
   "status": "running",
   "stopped": false,
@@ -139,6 +135,7 @@ exports[`cell reducer > can initialize stdin > second stdin requested 1`] = `
     {
       "channel": "stdin",
       "data": "what is your name?",
+      "id": 2,
       "mimetype": "text/plain",
       "response": "Marimo!",
       "timestamp": 2,
@@ -146,6 +143,7 @@ exports[`cell reducer > can initialize stdin > second stdin requested 1`] = `
     {
       "channel": "stdin",
       "data": "how old are you?",
+      "id": 3,
       "mimetype": "text/plain",
       "timestamp": 3,
     },
@@ -163,7 +161,6 @@ exports[`cell reducer > can initialize stdin > second stdin requested 1`] = `
   "output": null,
   "runElapsedTimeMs": null,
   "runStartTimestamp": 20,
-  "serializedEditorState": null,
   "staleInputs": false,
   "status": "running",
   "stopped": false,
@@ -189,6 +186,7 @@ exports[`cell reducer > can initialize stdin > stdin answered 1`] = `
     {
       "channel": "stdin",
       "data": "what is your name?",
+      "id": 2,
       "mimetype": "text/plain",
       "response": "Marimo!",
       "timestamp": 2,
@@ -207,7 +205,6 @@ exports[`cell reducer > can initialize stdin > stdin answered 1`] = `
   "output": null,
   "runElapsedTimeMs": null,
   "runStartTimestamp": 20,
-  "serializedEditorState": null,
   "staleInputs": false,
   "status": "running",
   "stopped": false,
@@ -233,6 +230,7 @@ exports[`cell reducer > can initialize stdin > stdin requested 1`] = `
     {
       "channel": "stdin",
       "data": "what is your name?",
+      "id": 2,
       "mimetype": "text/plain",
       "timestamp": 2,
     },
@@ -250,7 +248,6 @@ exports[`cell reducer > can initialize stdin > stdin requested 1`] = `
   "output": null,
   "runElapsedTimeMs": null,
   "runStartTimestamp": 20,
-  "serializedEditorState": null,
   "staleInputs": false,
   "status": "running",
   "stopped": false,
@@ -287,7 +284,6 @@ exports[`cell reducer > can initialize stdin > stdout received 1`] = `
   "output": null,
   "runElapsedTimeMs": null,
   "runStartTimestamp": 20,
-  "serializedEditorState": null,
   "staleInputs": false,
   "status": "running",
   "stopped": false,
@@ -324,7 +320,6 @@ exports[`cell reducer > can initialize stdin > stdout verified 1`] = `
   "output": null,
   "runElapsedTimeMs": null,
   "runStartTimestamp": 20,
-  "serializedEditorState": null,
   "staleInputs": false,
   "status": "running",
   "stopped": false,
@@ -354,7 +349,6 @@ exports[`cell reducer > can receive console when the cell is idle and will clear
   "output": null,
   "runElapsedTimeMs": null,
   "runStartTimestamp": 20,
-  "serializedEditorState": null,
   "staleInputs": false,
   "status": "running",
   "stopped": false,
@@ -391,7 +385,6 @@ exports[`cell reducer > can receive console when the cell is idle and will clear
   "output": null,
   "runElapsedTimeMs": 2000,
   "runStartTimestamp": null,
-  "serializedEditorState": null,
   "staleInputs": false,
   "status": "idle",
   "stopped": false,
@@ -421,7 +414,6 @@ exports[`cell reducer > can receive console when the cell is idle and will clear
   "output": null,
   "runElapsedTimeMs": null,
   "runStartTimestamp": null,
-  "serializedEditorState": null,
   "staleInputs": false,
   "status": "idle",
   "stopped": false,
@@ -458,7 +450,6 @@ exports[`cell reducer > can receive console when the cell is idle and will clear
   "output": null,
   "runElapsedTimeMs": null,
   "runStartTimestamp": null,
-  "serializedEditorState": null,
   "staleInputs": false,
   "status": "queued",
   "stopped": false,
@@ -495,7 +486,6 @@ exports[`cell reducer > can receive console when the cell is idle and will clear
   "output": null,
   "runElapsedTimeMs": null,
   "runStartTimestamp": null,
-  "serializedEditorState": null,
   "staleInputs": false,
   "status": "queued",
   "stopped": false,
@@ -525,7 +515,6 @@ exports[`cell reducer > can run a stale cell > queued message received 1`] = `
   "output": null,
   "runElapsedTimeMs": null,
   "runStartTimestamp": null,
-  "serializedEditorState": null,
   "staleInputs": false,
   "status": "queued",
   "stopped": false,
@@ -555,7 +544,6 @@ exports[`cell reducer > can run a stale cell > stale inputs received 1`] = `
   "output": null,
   "runElapsedTimeMs": null,
   "runStartTimestamp": null,
-  "serializedEditorState": null,
   "staleInputs": true,
   "status": "queued",
   "stopped": false,
@@ -596,7 +584,6 @@ exports[`cell reducer > can run a stopped cell > ancestor stopped 1`] = `
   },
   "runElapsedTimeMs": null,
   "runStartTimestamp": null,
-  "serializedEditorState": null,
   "staleInputs": false,
   "status": "idle",
   "stopped": true,
@@ -626,7 +613,6 @@ exports[`cell reducer > can run a stopped cell > idle message received 1`] = `
   "output": null,
   "runElapsedTimeMs": null,
   "runStartTimestamp": null,
-  "serializedEditorState": null,
   "staleInputs": false,
   "status": "idle",
   "stopped": false,
@@ -667,7 +653,6 @@ exports[`cell reducer > can run a stopped cell > queued after stop 1`] = `
   },
   "runElapsedTimeMs": null,
   "runStartTimestamp": null,
-  "serializedEditorState": null,
   "staleInputs": false,
   "status": "queued",
   "stopped": true,
@@ -697,7 +682,6 @@ exports[`cell reducer > can run a stopped cell > queued message received 1`] = `
   "output": null,
   "runElapsedTimeMs": null,
   "runStartTimestamp": null,
-  "serializedEditorState": null,
   "staleInputs": false,
   "status": "queued",
   "stopped": false,
@@ -727,7 +711,6 @@ exports[`cell reducer > can run a stopped cell > running after stop 1`] = `
   "output": null,
   "runElapsedTimeMs": null,
   "runStartTimestamp": 40,
-  "serializedEditorState": null,
   "staleInputs": false,
   "status": "running",
   "stopped": false,
@@ -770,7 +753,6 @@ import numpy",
   },
   "runElapsedTimeMs": 13000,
   "runStartTimestamp": null,
-  "serializedEditorState": null,
   "staleInputs": false,
   "status": "idle",
   "stopped": false,
@@ -800,7 +782,6 @@ exports[`cell reducer > can run cell and receive cell messages > code edited 1`]
   "output": null,
   "runElapsedTimeMs": null,
   "runStartTimestamp": null,
-  "serializedEditorState": null,
   "staleInputs": false,
   "status": "idle",
   "stopped": false,
@@ -843,7 +824,6 @@ import numpy",
   },
   "runElapsedTimeMs": 13000,
   "runStartTimestamp": null,
-  "serializedEditorState": null,
   "staleInputs": false,
   "status": "idle",
   "stopped": false,
@@ -885,7 +865,6 @@ exports[`cell reducer > can run cell and receive cell messages > code reverted 1
   },
   "runElapsedTimeMs": 13000,
   "runStartTimestamp": null,
-  "serializedEditorState": null,
   "staleInputs": false,
   "status": "idle",
   "stopped": false,
@@ -922,7 +901,6 @@ exports[`cell reducer > can run cell and receive cell messages > console message
   "output": null,
   "runElapsedTimeMs": null,
   "runStartTimestamp": 20,
-  "serializedEditorState": null,
   "staleInputs": false,
   "status": "running",
   "stopped": false,
@@ -972,7 +950,6 @@ import numpy",
   },
   "runElapsedTimeMs": 11000,
   "runStartTimestamp": null,
-  "serializedEditorState": null,
   "staleInputs": false,
   "status": "idle",
   "stopped": false,
@@ -1002,7 +979,6 @@ exports[`cell reducer > can run cell and receive cell messages > initial state 1
   "output": null,
   "runElapsedTimeMs": null,
   "runStartTimestamp": null,
-  "serializedEditorState": null,
   "staleInputs": false,
   "status": "idle",
   "stopped": false,
@@ -1050,7 +1026,6 @@ import numpy",
   },
   "runElapsedTimeMs": 11000,
   "runStartTimestamp": null,
-  "serializedEditorState": null,
   "staleInputs": false,
   "status": "idle",
   "stopped": false,
@@ -1092,7 +1067,6 @@ exports[`cell reducer > can run cell and receive cell messages > output received
   },
   "runElapsedTimeMs": 13000,
   "runStartTimestamp": null,
-  "serializedEditorState": null,
   "staleInputs": false,
   "status": "idle",
   "stopped": false,
@@ -1122,7 +1096,6 @@ exports[`cell reducer > can run cell and receive cell messages > prepared for ru
   "output": null,
   "runElapsedTimeMs": null,
   "runStartTimestamp": null,
-  "serializedEditorState": null,
   "staleInputs": false,
   "status": "queued",
   "stopped": false,
@@ -1152,7 +1125,6 @@ exports[`cell reducer > can run cell and receive cell messages > queued message
   "output": null,
   "runElapsedTimeMs": null,
   "runStartTimestamp": null,
-  "serializedEditorState": null,
   "staleInputs": false,
   "status": "queued",
   "stopped": false,
@@ -1182,7 +1154,6 @@ exports[`cell reducer > can run cell and receive cell messages > running message
   "output": null,
   "runElapsedTimeMs": null,
   "runStartTimestamp": 20,
-  "serializedEditorState": null,
   "staleInputs": false,
   "status": "running",
   "stopped": false,
@@ -1212,7 +1183,6 @@ exports[`cell reducer > errors reset status to idle > code edited 1`] = `
   "output": null,
   "runElapsedTimeMs": null,
   "runStartTimestamp": null,
-  "serializedEditorState": null,
   "staleInputs": false,
   "status": "idle",
   "stopped": false,
@@ -1253,7 +1223,6 @@ exports[`cell reducer > errors reset status to idle > error response received 1`
   },
   "runElapsedTimeMs": null,
   "runStartTimestamp": null,
-  "serializedEditorState": null,
   "staleInputs": false,
   "status": "idle",
   "stopped": false,
@@ -1283,7 +1252,6 @@ exports[`cell reducer > errors reset status to idle > initial state 1`] = `
   "output": null,
   "runElapsedTimeMs": null,
   "runStartTimestamp": null,
-  "serializedEditorState": null,
   "staleInputs": false,
   "status": "idle",
   "stopped": false,
@@ -1313,7 +1281,6 @@ exports[`cell reducer > errors reset status to idle > prepared for run 1`] = `
   "output": null,
   "runElapsedTimeMs": null,
   "runStartTimestamp": null,
-  "serializedEditorState": null,
   "staleInputs": false,
   "status": "queued",
   "stopped": false,