@marimo-team/islands 0.25.3-dev26 → 0.25.3-dev262

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/cell.ts

@@ -0,0 +1,54 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+import { notebookAtom } from "@/core/cells/cells";
+import { displayCellName } from "@/core/cells/names";
+import type { JotaiStore } from "@/core/state/jotai";
+import { AIContextProvider, type AIContextItem } from "../registry";
+import { contextToXml } from "../utils";
+
+interface CellContextItem extends AIContextItem {
+  type: "cell";
+  data: { cellId: string; code: string };
+}
+
+export class CellContextProvider extends AIContextProvider<CellContextItem> {
+  public readonly title = "Cells";
+  public readonly mentionPrefix = "@";
+  public readonly contextType = "cell";
+
+  private readonly store: JotaiStore;
+
+  public constructor(store: JotaiStore) {
+    super();
+    this.store = store;
+  }
+
+  public getItems(): CellContextItem[] {
+    const notebook = this.store.get(notebookAtom);
+    return notebook.cellIds.map((cellId, index) => {
+      const cell = notebook.cellData[cellId];
+      return {
+        uri: this.asURI(cellId),
+        name: displayCellName(cell.name, index),
+        type: this.contextType,
+        description: "Cell code",
+        data: { cellId, code: cell.code },
+      };
+    });
+  }
+
+  public formatCompletion(item: CellContextItem) {
+    return {
+      ...this.createBasicCompletion(item),
+      label: `${this.mentionPrefix}${item.name}`,
+    };
+  }
+
+  public formatContext(item: CellContextItem) {
+    return contextToXml({
+      type: "cell",
+      data: { cellId: item.data.cellId, name: item.name },
+      details: item.data.code || "This cell is empty.",
+    });
+  }
+}

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

@@ -1,33 +1,19 @@
 /* Copyright 2026 Marimo. All rights reserved. */
 
-import { atom, useAtomValue } from "jotai";
+import { useAtomValue } from "jotai";
 import { selectAtom } from "jotai/utils";
-import { useMemo, useRef } from "react";
+import { useMemo } from "react";
 import { useDeleteCellCallback } from "@/components/editor/cell/useDeleteCell";
 import { CellId } from "@/core/cells/ids";
 import { createReducerAndAtoms } from "@/utils/createReducer";
 import { Logger } from "@/utils/Logger";
-import {
-  type CompletionDataPart,
-  NOTEBOOK_CELLS_COMPLETION_DATA_TYPE,
-  type GeneratedCell,
-  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 { 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";
 
-/**
- * Pending AI edits shared by completion, chat, and agent workflows.
- * Generation ownership keeps one workflow from accepting or discarding another's
- * staged cells.
- */
+/** Pending AI edits shared by chat and agent workflows. */
 
 export type Edit =
   | { type: Extract<EditType, "update_cell">; previousCode: string }
@@ -36,25 +22,6 @@ export type Edit =
 
 export type StagedAICells = Map<CellId, Edit>;
 
-type StagedGeneration =
-  | {
-      id: symbol;
-      status: "in_progress";
-      cellIds: readonly CellId[];
-    }
-  | { id: symbol; status: "complete" };
-
-interface OwnedStagedGeneration {
-  id: symbol;
-  status: "in_progress" | "complete";
-  reconciler: StagedCellReconciler;
-}
-
-const stagedGenerationAtom = atom<StagedGeneration | null>(null);
-export const stagedGenerationInProgressAtom = atom(
-  (get) => get(stagedGenerationAtom)?.status === "in_progress",
-);
-
 const initialState = (): StagedAICells => {
   return new Map();
 };
@@ -98,18 +65,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 +84,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 });
     }
   };
 
@@ -144,255 +113,9 @@ 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 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.
-  useAtomValue(stagedGenerationAtom, { store });
-  useAtomValue(stagedAICellsAtom, { store });
-
-  const discardStagedCellIds = (cellIds: readonly CellId[]) => {
-    const stagedAICells = store.get(stagedAICellsAtom);
-    for (const cellId of cellIds.toReversed()) {
-      if (stagedAICells.has(cellId)) {
-        stagedCellActions.deleteStagedCell(cellId);
-      }
-    }
-  };
-
-  const beginStagedCellGeneration = () => {
-    const activeGeneration = store.get(stagedGenerationAtom);
-    if (activeGeneration?.status === "in_progress") {
-      discardStagedCellIds(activeGeneration.cellIds);
-    }
-    ownedGeneration.current?.reconciler.discard(store.get(stagedAICellsAtom));
-    const generationId = Symbol("staged-cell-generation");
-    ownedGeneration.current = {
-      id: generationId,
-      status: "in_progress",
-      reconciler: new StagedCellReconciler({
-        ...stagedCellActions,
-        createNewCell,
-      }),
-    };
-    store.set(stagedGenerationAtom, {
-      id: generationId,
-      status: "in_progress",
-      cellIds: [],
-    });
-  };
-
-  const finishStagedCellGeneration = (successful: boolean) => {
-    const generation = ownedGeneration.current;
-    if (generation?.status !== "in_progress") {
-      return;
-    }
-
-    const stagedGeneration = store.get(stagedGenerationAtom);
-    const ownsActiveGeneration =
-      stagedGeneration?.id === generation.id &&
-      stagedGeneration.status === "in_progress";
-    if (!ownsActiveGeneration || !successful) {
-      generation.reconciler.discard(store.get(stagedAICellsAtom));
-      if (ownsActiveGeneration) {
-        store.set(stagedGenerationAtom, null);
-      }
-      ownedGeneration.current = null;
-      return;
-    }
-
-    generation.status = "complete";
-    store.set(stagedGenerationAtom, {
-      id: generation.id,
-      status: "complete",
-    });
-  };
-
-  const hasOwnedStagedCells = () => {
-    const generation = ownedGeneration.current;
-    const stagedGeneration = store.get(stagedGenerationAtom);
-    return (
-      generation !== null &&
-      generation.status === "complete" &&
-      stagedGeneration?.status === "complete" &&
-      stagedGeneration.id === generation.id &&
-      generation.reconciler.hasStagedCells(store.get(stagedAICellsAtom))
-    );
-  };
-
-  const acceptOwnedStagedCells = () => {
-    if (!hasOwnedStagedCells()) {
-      return false;
-    }
-
-    ownedGeneration.current?.reconciler.accept();
-    store.set(stagedGenerationAtom, null);
-    ownedGeneration.current = null;
-    return true;
-  };
-
-  const discardOwnedStagedCells = () => {
-    const generation = ownedGeneration.current;
-    const stagedGeneration = store.get(stagedGenerationAtom);
-    if (generation === null || stagedGeneration?.id !== generation.id) {
-      return false;
-    }
-
-    generation.reconciler.discard(store.get(stagedAICellsAtom));
-    store.set(stagedGenerationAtom, null);
-    ownedGeneration.current = null;
-    return true;
-  };
-
-  const onData = (part: CompletionDataPart) => {
-    if (part.type !== NOTEBOOK_CELLS_COMPLETION_DATA_TYPE) {
-      return;
-    }
-    const generation = ownedGeneration.current;
-    const stagedGeneration = store.get(stagedGenerationAtom);
-    if (
-      generation?.status !== "in_progress" ||
-      stagedGeneration?.id !== generation.id ||
-      stagedGeneration.status !== "in_progress"
-    ) {
-      return;
-    }
-    const completion = notebookCellsCompletionSchema.parse(part.data);
-    generation.reconciler.reconcile(completion.cells);
-    store.set(stagedGenerationAtom, {
-      ...stagedGeneration,
-      cellIds: generation.reconciler.stagedCellIds(),
-    });
-  };
-
-  return {
-    beginStagedCellGeneration,
-    finishStagedCellGeneration,
-    hasOwnedStagedCells,
-    acceptOwnedStagedCells,
-    discardOwnedStagedCells,
-    onData,
-  };
-}
-
 export { stagedAICellsAtom };
 export const visibleForTesting = {
-  stagedGenerationAtom,
   createActions,
   reducer,
   initialState,
 };
-
-interface CreatedCell {
-  cellId: CellId;
-  code: string;
-}
-
-interface StagedCellReconcilerActions {
-  createStagedCell: (code: string) => CellId;
-  updateStagedCell: (opts: UpdateStagedCellAction) => void;
-  deleteStagedCell: (cellId: CellId) => void;
-  removeStagedCell: (cellId: CellId) => void;
-  addStagedCell: (payload: { cellId: CellId; edit: Edit }) => void;
-  createNewCell: (opts: CreateNewCellAction) => void;
-}
-
-class StagedCellReconciler {
-  private createdCells: CreatedCell[] = [];
-  private readonly actions: StagedCellReconcilerActions;
-  private marimoImportCellId: CellId | null = null;
-
-  public constructor(actions: StagedCellReconcilerActions) {
-    this.actions = actions;
-  }
-
-  public reconcile(completionCells: GeneratedCell[]) {
-    for (const [idx, cell] of completionCells.entries()) {
-      if (idx < this.createdCells.length) {
-        const existingCell = this.createdCells[idx];
-        const codeChanged = existingCell.code !== cell.code;
-        this.createdCells[idx] = { ...existingCell, code: cell.code };
-        if (!codeChanged) {
-          continue;
-        }
-        this.actions.updateStagedCell({
-          cellId: existingCell.cellId,
-          code: cell.code,
-        });
-      } else {
-        const newCellId = this.actions.createStagedCell(cell.code);
-        this.createdCells.push({ cellId: newCellId, code: cell.code });
-      }
-    }
-
-    const removedCells = this.createdCells.splice(completionCells.length);
-    for (const { cellId } of removedCells.toReversed()) {
-      this.actions.deleteStagedCell(cellId);
-    }
-
-    this.syncMarimoImport(completionCells);
-  }
-
-  public discard(stagedCells: StagedAICells) {
-    for (const { cellId } of this.createdCells.toReversed()) {
-      if (stagedCells.has(cellId)) {
-        this.actions.deleteStagedCell(cellId);
-      }
-    }
-    this.createdCells = [];
-
-    if (this.marimoImportCellId && stagedCells.has(this.marimoImportCellId)) {
-      this.actions.deleteStagedCell(this.marimoImportCellId);
-    }
-    this.marimoImportCellId = null;
-  }
-
-  public accept() {
-    for (const cellId of this.stagedCellIds()) {
-      this.actions.removeStagedCell(cellId);
-    }
-  }
-
-  public hasStagedCells(stagedCells: StagedAICells) {
-    return this.stagedCellIds().some((cellId) => stagedCells.has(cellId));
-  }
-
-  public stagedCellIds() {
-    const cellIds = this.createdCells.map(({ cellId }) => cellId);
-    if (this.marimoImportCellId) {
-      cellIds.push(this.marimoImportCellId);
-    }
-    return cellIds;
-  }
-
-  /** Keep the generated marimo import consistent with the latest snapshot. */
-  private syncMarimoImport(completionCells: GeneratedCell[]) {
-    const requiresMarimo = completionCells.some(
-      (cell) => cell.language !== "python",
-    );
-    if (!requiresMarimo) {
-      if (this.marimoImportCellId) {
-        this.actions.deleteStagedCell(this.marimoImportCellId);
-        this.marimoImportCellId = null;
-      }
-      return;
-    }
-    if (this.marimoImportCellId) {
-      return;
-    }
-
-    const cellId = maybeAddMarimoImport({
-      autoInstantiate: false,
-      createNewCell: this.actions.createNewCell,
-      fromCellId: this.createdCells[0]?.cellId,
-      before: true,
-    });
-    if (cellId) {
-      this.actions.addStagedCell({ cellId, edit: { type: "add_cell" } });
-      this.marimoImportCellId = cellId;
-    }
-  }
-}

package/src/core/ai/state.ts

@@ -31,10 +31,13 @@ export const aiCompletionCellAtom = atom<AiCompletionCell | null>(null);
 export interface PendingAiPrompt {
   prompt: string;
   submit: boolean;
+  append?: boolean;
 }
 
 export const pendingAiPromptAtom = atom<PendingAiPrompt | null>(null);
 
+export const chatGenerationInProgressAtom = atom(false);
+
 const CHAT_OPTIONS_KEY = "marimo:ai:chatOptions";
 export const chatOptionsAtom = atomWithStorage<Required<ChatOptions>>(
   CHAT_OPTIONS_KEY,

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