@marimo-team/islands 0.25.3-dev23 → 0.25.3-dev231

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

@@ -7,7 +7,7 @@ import { beforeEach, describe, expect, it, vi } from "vitest";
 import { cellId } from "@/__tests__/branded";
 import { CellId } from "@/core/cells/ids";
 import { getCellEditorView } from "../../cells/cells";
-import { updateEditorCodeFromPython } from "../../codemirror/language/utils";
+import { prepareCellCodeUpdate } from "../../codemirror/language/cell-code";
 import {
   stagedAICellsAtom,
   stagedGenerationInProgressAtom,
@@ -21,8 +21,8 @@ const { stagedGenerationAtom, createActions, reducer, initialState } =
   visibleForTesting;
 
 // Mock the dependencies
-const mockCreateNewCell = vi.fn();
-const mockUpdateCellCode = vi.fn();
+const mockCreateCell = vi.fn();
+const mockSetCellCode = vi.fn();
 const mockDeleteCellCallback = vi.fn();
 
 // Mock cell handle with editor view
@@ -34,17 +34,16 @@ const mockCellHandle = {
   },
 };
 
+vi.mock("../../cells/actions", () => ({
+  useCellActions: () => ({
+    createCell: mockCreateCell,
+    setCellCode: mockSetCellCode,
+  }),
+}));
 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),
-    })),
+    notebookAtom: atom({ cellData: {}, cellIds: [] }),
     getCellEditorView: vi.fn(() => mockCellHandle.current.editorViewOrNull),
   };
 });
@@ -57,8 +56,8 @@ vi.mock("@/core/network/requests", () => ({
   getRequestClient: () => ({ sendRun: vi.fn() }),
 }));
 
-vi.mock("../../codemirror/language/utils", () => ({
-  updateEditorCodeFromPython: vi.fn(),
+vi.mock("../../codemirror/language/cell-code", () => ({
+  prepareCellCodeUpdate: vi.fn(() => ({ effects: [] })),
 }));
 
 // Mock CellId.create
@@ -261,11 +260,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 +327,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,8 +398,9 @@ describe("staged cell generation", () => {
     result.current.beginStagedCellGeneration();
 
     // No cell or cell update should have been called
-    expect(mockCreateNewCell).not.toHaveBeenCalled();
-    expect(updateEditorCodeFromPython).not.toHaveBeenCalled();
+    expect(mockCreateCell).not.toHaveBeenCalled();
+    expect(prepareCellCodeUpdate).not.toHaveBeenCalled();
+    expect(mockSetCellCode).not.toHaveBeenCalled();
     expect(store.get(stagedGenerationInProgressAtom)).toBe(true);
   });
 
@@ -610,11 +608,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 +632,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,13 +679,17 @@ 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(updateEditorCodeFromPython).toHaveBeenCalledTimes(2);
+    expect(prepareCellCodeUpdate).toHaveBeenCalledTimes(2);
+    expect(
+      mockCellHandle.current.editorViewOrNull.dispatch,
+    ).toHaveBeenCalledTimes(2);
+    expect(mockSetCellCode).not.toHaveBeenCalled();
   });
 
-  it("should update notebook state when the cell editor is not mounted", () => {
+  it("writes through the document when the cell editor is not mounted", () => {
     const { result } = renderGenerationHook();
     result.current.beginStagedCellGeneration();
 
@@ -709,10 +709,9 @@ describe("staged cell generation", () => {
       },
     });
 
-    expect(mockUpdateCellCode).toHaveBeenCalledWith({
+    expect(mockSetCellCode).toHaveBeenCalledWith({
       cellId: cellId("unmounted-cell"),
       code: "value = 1",
-      formattingChange: false,
     });
   });
 

package/src/core/ai/context/providers/__tests__/cell-output.test.ts

@@ -34,9 +34,7 @@ describe("CellOutputContextProvider", () => {
   beforeEach(() => {
     // Create a basic mock notebook state
     mockNotebook = {
-      cellIds: {
-        inOrderIds: ["cell1", "cell2", "cell3"],
-      },
+      cellIds: ["cell1", "cell2", "cell3"],
       cellData: {
         cell1: {
           id: "cell1",
@@ -207,7 +205,7 @@ describe("Cell output utility functions", () => {
 
     beforeEach(() => {
       const mockNotebook = {
-        cellIds: { inOrderIds: [] },
+        cellIds: [],
         cellData: {},
         cellRuntime: {},
       } as unknown as NotebookState;
@@ -226,7 +224,7 @@ describe("Cell output utility functions", () => {
 
       for (const testCase of testCases) {
         mockStore = createMockStore({
-          cellIds: { inOrderIds: ["test"] },
+          cellIds: ["test"],
           cellData: {
             test: { id: "test", name: "", code: "" },
           },
@@ -261,7 +259,7 @@ describe("Cell output utility functions", () => {
 
       for (const testCase of testCases) {
         mockStore = createMockStore({
-          cellIds: { inOrderIds: ["test"] },
+          cellIds: ["test"],
           cellData: {
             test: { id: "test", name: "", code: "" },
           },
@@ -295,7 +293,7 @@ describe("Cell output utility functions", () => {
 
       for (const testCase of testCases) {
         mockStore = createMockStore({
-          cellIds: { inOrderIds: ["test"] },
+          cellIds: ["test"],
           cellData: {
             test: { id: "test", name: "", code: "" },
           },
@@ -338,7 +336,7 @@ describe("Cell output utility functions", () => {
       } as unknown as Document;
 
       const mockNotebook = {
-        cellIds: { inOrderIds: [] },
+        cellIds: [],
         cellData: {},
         cellRuntime: {},
       } as unknown as NotebookState;
@@ -369,7 +367,7 @@ describe("Cell output utility functions", () => {
       ).mockReturnValue(mockDiv);
 
       mockStore = createMockStore({
-        cellIds: { inOrderIds: ["test"] },
+        cellIds: ["test"],
         cellData: {
           test: { id: "test", name: "", code: "" },
         },
@@ -395,9 +393,7 @@ describe("Cell output utility functions", () => {
 
     beforeEach(() => {
       mockNotebook = {
-        cellIds: {
-          inOrderIds: ["cell1", "cell2"],
-        },
+        cellIds: ["cell1", "cell2"],
         cellData: {
           cell1: {
             id: "cell1",

package/src/core/ai/context/providers/cell-output.ts

@@ -93,7 +93,7 @@ export class CellOutputContextProvider extends AIContextProvider<CellOutputConte
     const notebook = this.store.get(notebookAtom);
     const items: CellOutputContextItem[] = [];
 
-    for (const cellId of notebook.cellIds.inOrderIds) {
+    for (const cellId of notebook.cellIds) {
       const cellContextData = getCellContextData(cellId, notebook);
       const cellOutput = cellContextData.cellOutput;
 
@@ -358,7 +358,7 @@ export function getCellContextData(
   const cellRuntime = notebook.cellRuntime[cellId];
 
   const cellData = notebook.cellData[cellId];
-  const cellIndex = notebook.cellIds.inOrderIds.indexOf(cellId);
+  const cellIndex = notebook.cellIds.indexOf(cellId);
   const cellName = displayCellName(cellData.name, cellIndex);
 
   let consoleOutputs: BaseOutput[] | undefined;

package/src/core/ai/context/providers/datasource.ts

@@ -98,8 +98,12 @@ export class DatasourceContextProvider extends AIContextProvider<DatasourceConte
   public formatContext(item: DatasourceContextItem): string {
     const data = item.data;
     // Remove certain fields that are not needed in the context
-    const { name, display_name, source, ...filteredDatasource } =
-      data.connection;
+    const {
+      name,
+      display_name: _displayName,
+      source: _source,
+      ...filteredDatasource
+    } = data.connection;
 
     let datasource = filteredDatasource;
     const isInternalEngine = INTERNAL_SQL_ENGINES.has(name as ConnectionName);

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

@@ -14,12 +14,9 @@ import {
   notebookCellsCompletionSchema,
 } from "./completion-output";
 import { maybeAddMarimoImport } from "../cells/add-missing-import";
-import {
-  type CreateNewCellAction,
-  getCellEditorView,
-  useCellActions,
-} from "../cells/cells";
-import { updateEditorCodeFromPython } from "../codemirror/language/utils";
+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 { setCellCode, 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;
   };
 
@@ -122,11 +114,18 @@ 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) {
-      updateEditorCodeFromPython(editorView, code);
+      const update = prepareCellCodeUpdate(editorView.state, code, {
+        initialize: true,
+      });
+      if (update) {
+        editorView.dispatch(update);
+      }
     } else {
-      updateCellCode({ cellId, code, formattingChange: false });
+      setCellCode({ cellId, code });
     }
   };
 
@@ -147,7 +146,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.
@@ -175,7 +174,7 @@ export function useStagedCellGeneration(store: JotaiStore) {
       status: "in_progress",
       reconciler: new StagedCellReconciler({
         ...stagedCellActions,
-        createNewCell,
+        createCell,
       }),
     };
     store.set(stagedGenerationAtom, {
@@ -297,7 +296,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 {
@@ -386,7 +385,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

@@ -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> {