@marimo-team/islands 0.25.3-dev34 → 0.25.3-dev348

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

@@ -7,6 +7,7 @@ import {
 import { getRequestClient } from "@/core/network/requests";
 import type { JotaiStore } from "@/core/state/jotai";
 import { variablesAtom } from "@/core/variables/state";
+import { CellContextProvider } from "./providers/cell";
 import { CellOutputContextProvider } from "./providers/cell-output";
 import { DatasourceContextProvider } from "./providers/datasource";
 import { ErrorContextProvider } from "./providers/error";
@@ -25,6 +26,7 @@ export function getAIContextRegistry(store: JotaiStore) {
     .register(new VariableContextProvider(variables, tablesMap))
     .register(new ErrorContextProvider(store))
     .register(new CellOutputContextProvider(store))
+    .register(new CellContextProvider(store))
     .register(
       new DatasourceContextProvider(datasource.connectionsMap, tablesMap),
     );

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

@@ -0,0 +1,49 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+import { createStore } from "jotai";
+import { expect, it } from "vitest";
+import { MockNotebook } from "@/__mocks__/notebook";
+import { notebookAtom } from "@/core/cells/cells";
+import { CellId } from "@/core/cells/ids";
+import { getAIContextRegistry } from "../../context";
+import { CellContextProvider } from "../cell";
+
+it("resolves empty cells and reads their latest code when preparing context", () => {
+  const store = createStore();
+  const cellId = CellId.create();
+  const provider = new CellContextProvider(store);
+  store.set(
+    notebookAtom,
+    MockNotebook.notebookState({ cellData: { [cellId]: { code: "" } } }),
+  );
+  const uri = provider.asURI(cellId);
+  expect(getAIContextRegistry(store).formatContextForAI([uri])).toContain(
+    "This cell is empty.",
+  );
+  store.set(
+    notebookAtom,
+    MockNotebook.notebookState({ cellData: { [cellId]: { code: "x = 1" } } }),
+  );
+  expect(getAIContextRegistry(store).formatContextForAI([uri])).toContain(
+    "x = 1",
+  );
+  store.set(notebookAtom, MockNotebook.notebookState());
+  expect(getAIContextRegistry(store).formatContextForAI([uri])).toBe("");
+});
+
+it("completes cell names while inserting stable cell references", () => {
+  const store = createStore();
+  const cellId = CellId.create();
+  store.set(
+    notebookAtom,
+    MockNotebook.notebookState({
+      cellData: { [cellId]: { name: "my_analysis", code: "x = 1" } },
+    }),
+  );
+  const provider = new CellContextProvider(store);
+  expect(provider.formatCompletion(provider.getItems()[0])).toMatchObject({
+    label: "@my_analysis",
+    displayLabel: "my_analysis",
+    apply: `@cell://${cellId}`,
+  });
+});

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,17 @@
 /* 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 type { 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 type { JotaiStore } from "../state/jotai";
+import { getCellEditorView } from "../cells/cells";
+import { useCellActions } from "../cells/actions";
+import { prepareCellCodeUpdate } from "../codemirror/language/cell-code";
 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 +20,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();
 };
@@ -69,6 +34,18 @@ const {
     const { cellId, edit } = action;
     return new Map([...state, [cellId, edit]]);
   },
+  reconcileStagedCells: (state, cellIds: readonly string[]) => {
+    const existing = new Set(cellIds);
+    const removed = [...state.keys()].filter((id) => !existing.has(id));
+    if (removed.length === 0) {
+      return state;
+    }
+    const next = new Map(state);
+    for (const id of removed) {
+      next.delete(id);
+    }
+    return next;
+  },
   removeStagedCell: (state, cellId: CellId) => {
     if (!state.has(cellId)) {
       return state;
@@ -81,52 +58,37 @@ const {
 
 export { useStagedAICellsActions };
 
-export function useStagedAICell(cellId: CellId): Edit | undefined {
+export function useStagedAICell(cellId: CellId | undefined): Edit | undefined {
   const stagedCellAtom = useMemo(
     () =>
-      selectAtom(stagedAICellsAtom, (cells) => cells.get(cellId), Object.is),
+      selectAtom(
+        stagedAICellsAtom,
+        (cells) => (cellId === undefined ? undefined : cells.get(cellId)),
+        Object.is,
+      ),
     [cellId],
   );
   return useAtomValue(stagedCellAtom);
 }
 
-interface UpdateStagedCellAction {
-  cellId: CellId;
-  code: string;
-}
-
-/** Manage staged cells without assuming which AI workflow owns them. */
-export function useStagedCells(store: JotaiStore) {
-  const { addStagedCell, removeStagedCell } = useStagedAICellsActions();
-  const { createNewCell, updateCellCode } = useCellActions();
+/** Keep or discard staged cells, whichever AI workflow staged them. */
+export function useStagedCells() {
+  const { removeStagedCell } = useStagedAICellsActions();
+  const { setCellCode } = 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,
-    });
-    return newCellId;
-  };
-
-  const updateStagedCell = (opts: UpdateStagedCellAction) => {
-    const { cellId, code } = opts;
-    const stagedAICells = store.get(stagedAICellsAtom);
-
-    if (!stagedAICells.has(cellId)) {
-      Logger.error("Staged cell not found", { cellId });
-      return;
-    }
-
+  // Only a mounted editor infers the cell's language; otherwise use the document.
+  const writeCellCode = (cellId: CellId, code: string) => {
     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 });
     }
   };
 
@@ -135,264 +97,35 @@ export function useStagedCells(store: JotaiStore) {
     deleteCellCallback({ cellId });
   };
 
-  return {
-    createStagedCell,
-    updateStagedCell,
-    addStagedCell,
-    removeStagedCell,
-    deleteStagedCell,
-  };
-}
-
-/** 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;
+  const acceptStagedCell = (cellId: CellId, edit: Edit) => {
+    if (edit.type === "delete_cell") {
+      deleteStagedCell(cellId);
+    } else {
+      removeStagedCell(cellId);
     }
-
-    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 rejectStagedCell = (cellId: CellId, edit: Edit) => {
+    switch (edit.type) {
+      case "update_cell":
+        writeCellCode(cellId, edit.previousCode);
+        removeStagedCell(cellId);
+        break;
+      case "add_cell":
+        deleteStagedCell(cellId);
+        break;
+      case "delete_cell":
+        removeStagedCell(cellId);
+        break;
     }
-    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,
-  };
+  return { acceptStagedCell, rejectStagedCell };
 }
 
 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,