@marimo-team/islands 0.25.3-dev250 → 0.25.3-dev252

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

@@ -6,19 +6,14 @@ import { createElement, type ReactNode } from "react";
 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 { prepareCellCodeUpdate } from "../../codemirror/language/cell-code";
 import {
   stagedAICellsAtom,
-  stagedGenerationInProgressAtom,
   useStagedAICell,
-  useStagedCellGeneration,
   useStagedCells,
   visibleForTesting,
 } from "../staged-cells";
 
-const { stagedGenerationAtom, createActions, reducer, initialState } =
-  visibleForTesting;
+const { createActions, reducer, initialState } = visibleForTesting;
 
 // Mock the dependencies
 const mockCreateCell = vi.fn();
@@ -83,7 +78,6 @@ describe("staged-cells", () => {
 
     // Reset the atom state
     store.set(stagedAICellsAtom, new Map());
-    store.set(stagedGenerationAtom, null);
   });
 
   describe("reducer and actions", () => {
@@ -379,463 +373,3 @@ describe("staged-cells", () => {
     });
   });
 });
-
-describe("staged cell generation", () => {
-  let store: ReturnType<typeof getDefaultStore>;
-  const renderGenerationHook = () =>
-    renderHook(() => useStagedCellGeneration(store));
-
-  beforeEach(() => {
-    store = getDefaultStore();
-    vi.clearAllMocks();
-    vi.mocked(CellId.create).mockReset();
-    store.set(stagedAICellsAtom, new Map());
-    store.set(stagedGenerationAtom, null);
-  });
-
-  it("should begin generation", () => {
-    const { result } = renderGenerationHook();
-    result.current.beginStagedCellGeneration();
-
-    // No cell or cell update should have been called
-    expect(mockCreateCell).not.toHaveBeenCalled();
-    expect(prepareCellCodeUpdate).not.toHaveBeenCalled();
-    expect(mockSetCellCode).not.toHaveBeenCalled();
-    expect(store.get(stagedGenerationInProgressAtom)).toBe(true);
-  });
-
-  it("should mark generation complete after a successful finish", () => {
-    const { result } = renderGenerationHook();
-    result.current.beginStagedCellGeneration();
-    result.current.finishStagedCellGeneration(true);
-
-    expect(store.get(stagedGenerationInProgressAtom)).toBe(false);
-  });
-
-  it("should not couple unrelated staged-cell removal to generation state", () => {
-    const { result: stagedCells } = renderHook(() => useStagedCells(store));
-    const { result } = renderGenerationHook();
-    result.current.beginStagedCellGeneration();
-
-    const chatCellId = cellId("chat-edited-cell");
-    stagedCells.current.addStagedCell({
-      cellId: chatCellId,
-      edit: { type: "update_cell", previousCode: "before" },
-    });
-    stagedCells.current.removeStagedCell(chatCellId);
-
-    expect(store.get(stagedGenerationInProgressAtom)).toBe(true);
-  });
-
-  it("should only accept cells owned by the completed generation", () => {
-    const firstGeneration = renderGenerationHook();
-    firstGeneration.result.current.beginStagedCellGeneration();
-
-    vi.mocked(CellId.create).mockReturnValueOnce(cellId("superseded-cell"));
-    firstGeneration.result.current.onData({
-      type: "data-notebook-cells-completion",
-      data: {
-        cells: [{ language: "python", code: "old" }],
-      },
-    });
-    firstGeneration.result.current.finishStagedCellGeneration(true);
-
-    const activeGeneration = renderGenerationHook();
-    activeGeneration.result.current.beginStagedCellGeneration();
-    vi.mocked(CellId.create).mockReturnValueOnce(cellId("active-cell"));
-    activeGeneration.result.current.onData({
-      type: "data-notebook-cells-completion",
-      data: {
-        cells: [{ language: "python", code: "active" }],
-      },
-    });
-    activeGeneration.result.current.finishStagedCellGeneration(true);
-
-    expect(firstGeneration.result.current.hasOwnedStagedCells()).toBe(false);
-    expect(firstGeneration.result.current.acceptOwnedStagedCells()).toBe(false);
-    expect(store.get(stagedAICellsAtom)).toEqual(
-      new Map([
-        [cellId("superseded-cell"), { type: "add_cell" }],
-        [cellId("active-cell"), { type: "add_cell" }],
-      ]),
-    );
-
-    expect(activeGeneration.result.current.hasOwnedStagedCells()).toBe(true);
-    expect(activeGeneration.result.current.acceptOwnedStagedCells()).toBe(true);
-    expect(store.get(stagedAICellsAtom)).toEqual(
-      new Map([[cellId("superseded-cell"), { type: "add_cell" }]]),
-    );
-  });
-
-  it("should replace only the previous generation owned by this hook", () => {
-    const { result: stagedCells } = renderHook(() => useStagedCells(store));
-    const { result } = renderGenerationHook();
-    stagedCells.current.addStagedCell({
-      cellId: cellId("chat-edited-cell"),
-      edit: { type: "update_cell", previousCode: "before" },
-    });
-    result.current.beginStagedCellGeneration();
-
-    vi.mocked(CellId.create).mockReturnValueOnce(cellId("generated-cell"));
-    result.current.onData({
-      type: "data-notebook-cells-completion",
-      data: {
-        cells: [{ language: "python", code: "generated" }],
-      },
-    });
-    result.current.finishStagedCellGeneration(true);
-    mockDeleteCellCallback.mockClear();
-
-    result.current.beginStagedCellGeneration();
-
-    expect(mockDeleteCellCallback).toHaveBeenCalledExactlyOnceWith({
-      cellId: cellId("generated-cell"),
-    });
-    expect(store.get(stagedAICellsAtom)).toEqual(
-      new Map([
-        [
-          cellId("chat-edited-cell"),
-          { type: "update_cell", previousCode: "before" },
-        ],
-      ]),
-    );
-    expect(store.get(stagedGenerationInProgressAtom)).toBe(true);
-  });
-
-  it("should leave staged cells from another workflow untouched when accepting", () => {
-    const { result: stagedCells } = renderHook(() => useStagedCells(store));
-    const { result } = renderGenerationHook();
-    result.current.beginStagedCellGeneration();
-
-    vi.mocked(CellId.create).mockReturnValueOnce(cellId("generated-cell"));
-    result.current.onData({
-      type: "data-notebook-cells-completion",
-      data: {
-        cells: [{ language: "python", code: "generated" }],
-      },
-    });
-    result.current.finishStagedCellGeneration(true);
-    stagedCells.current.addStagedCell({
-      cellId: cellId("chat-edited-cell"),
-      edit: { type: "update_cell", previousCode: "before" },
-    });
-
-    expect(result.current.acceptOwnedStagedCells()).toBe(true);
-    expect(store.get(stagedAICellsAtom)).toEqual(
-      new Map([
-        [
-          cellId("chat-edited-cell"),
-          { type: "update_cell", previousCode: "before" },
-        ],
-      ]),
-    );
-  });
-
-  it("should only discard the generation's remaining staged cells", () => {
-    const { result: stagedCells } = renderHook(() => useStagedCells(store));
-    const { result } = renderGenerationHook();
-    result.current.beginStagedCellGeneration();
-
-    vi.mocked(CellId.create)
-      .mockReturnValueOnce(cellId("accepted-cell"))
-      .mockReturnValueOnce(cellId("pending-cell"));
-    result.current.onData({
-      type: "data-notebook-cells-completion",
-      data: {
-        cells: [
-          { language: "python", code: "accepted" },
-          { language: "python", code: "pending" },
-        ],
-      },
-    });
-    result.current.finishStagedCellGeneration(true);
-    stagedCells.current.removeStagedCell(cellId("accepted-cell"));
-    stagedCells.current.addStagedCell({
-      cellId: cellId("chat-edited-cell"),
-      edit: { type: "update_cell", previousCode: "before" },
-    });
-
-    expect(result.current.discardOwnedStagedCells()).toBe(true);
-    expect(mockDeleteCellCallback).not.toHaveBeenCalledWith({
-      cellId: cellId("accepted-cell"),
-    });
-    expect(mockDeleteCellCallback).toHaveBeenCalledWith({
-      cellId: cellId("pending-cell"),
-    });
-    expect(mockDeleteCellCallback).not.toHaveBeenCalledWith({
-      cellId: cellId("chat-edited-cell"),
-    });
-    expect(store.get(stagedAICellsAtom)).toEqual(
-      new Map([
-        [
-          cellId("chat-edited-cell"),
-          { type: "update_cell", previousCode: "before" },
-        ],
-      ]),
-    );
-  });
-
-  it.each([
-    { name: "no cells", cells: [] },
-    {
-      name: "an empty cell",
-      cells: [{ language: "python" as const, code: "" }],
-    },
-  ])("should reject $name", ({ cells }) => {
-    const { result } = renderGenerationHook();
-    result.current.beginStagedCellGeneration();
-
-    expect(() =>
-      result.current.onData({
-        type: "data-notebook-cells-completion",
-        data: { cells },
-      }),
-    ).toThrow(/cells/i);
-  });
-
-  it("should create cells from a validated completion", () => {
-    const { result } = renderGenerationHook();
-    result.current.beginStagedCellGeneration();
-
-    // Mock CellId.create to return a predictable ID
-    const mockCellId = cellId("mock-cell-id");
-    vi.mocked(CellId.create).mockReturnValue(mockCellId);
-
-    result.current.onData({
-      type: "data-notebook-cells-completion",
-      data: {
-        cells: [{ language: "python", code: "print('```')" }],
-      },
-    });
-
-    expect(mockCreateCell).toHaveBeenCalledWith({
-      code: "print('```')",
-      id: "mock-cell-id",
-    });
-  });
-
-  it("should create multiple cells", () => {
-    const { result } = renderGenerationHook();
-    result.current.beginStagedCellGeneration();
-
-    vi.mocked(CellId.create)
-      .mockReturnValueOnce(cellId("first-cell"))
-      .mockReturnValueOnce(cellId("second-cell"));
-
-    result.current.onData({
-      type: "data-notebook-cells-completion",
-      data: {
-        cells: [
-          { language: "python", code: "value = 1" },
-          { language: "python", code: "value" },
-        ],
-      },
-    });
-
-    expect(mockCreateCell).toHaveBeenCalledWith(
-      expect.objectContaining({
-        code: "value = 1",
-        id: "first-cell",
-      }),
-    );
-    expect(mockCreateCell).toHaveBeenCalledWith(
-      expect.objectContaining({
-        code: "value",
-        id: "second-cell",
-      }),
-    );
-  });
-
-  it("should apply cumulative snapshots without recreating cells", () => {
-    const { result } = renderGenerationHook();
-    result.current.beginStagedCellGeneration();
-
-    vi.mocked(CellId.create)
-      .mockReturnValueOnce(cellId("first-cell"))
-      .mockReturnValueOnce(cellId("second-cell"));
-
-    result.current.onData({
-      type: "data-notebook-cells-completion",
-      data: {
-        cells: [{ language: "python", code: "value" }],
-      },
-    });
-    result.current.onData({
-      type: "data-notebook-cells-completion",
-      data: {
-        cells: [
-          { language: "python", code: "value = 1" },
-          { language: "python", code: "value" },
-        ],
-      },
-    });
-    result.current.onData({
-      type: "data-notebook-cells-completion",
-      data: {
-        cells: [
-          { language: "python", code: "value = 1" },
-          { language: "python", code: "value + 1" },
-        ],
-      },
-    });
-
-    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(prepareCellCodeUpdate).toHaveBeenCalledTimes(2);
-    expect(
-      mockCellHandle.current.editorViewOrNull.dispatch,
-    ).toHaveBeenCalledTimes(2);
-    expect(mockSetCellCode).not.toHaveBeenCalled();
-  });
-
-  it("writes through the document when the cell editor is not mounted", () => {
-    const { result } = renderGenerationHook();
-    result.current.beginStagedCellGeneration();
-
-    vi.mocked(CellId.create).mockReturnValueOnce(cellId("unmounted-cell"));
-    result.current.onData({
-      type: "data-notebook-cells-completion",
-      data: {
-        cells: [{ language: "python", code: "value" }],
-      },
-    });
-
-    vi.mocked(getCellEditorView).mockReturnValueOnce(null);
-    result.current.onData({
-      type: "data-notebook-cells-completion",
-      data: {
-        cells: [{ language: "python", code: "value = 1" }],
-      },
-    });
-
-    expect(mockSetCellCode).toHaveBeenCalledWith({
-      cellId: cellId("unmounted-cell"),
-      code: "value = 1",
-    });
-  });
-
-  it("should remove trailing cells when a retry produces a shorter snapshot", () => {
-    const { result } = renderGenerationHook();
-    result.current.beginStagedCellGeneration();
-
-    vi.mocked(CellId.create)
-      .mockReturnValueOnce(cellId("first-cell"))
-      .mockReturnValueOnce(cellId("second-cell"));
-
-    result.current.onData({
-      type: "data-notebook-cells-completion",
-      data: {
-        cells: [
-          { language: "python", code: "value = 1" },
-          { language: "python", code: "value" },
-        ],
-      },
-    });
-    result.current.onData({
-      type: "data-notebook-cells-completion",
-      data: {
-        cells: [{ language: "python", code: "replacement = 2" }],
-      },
-    });
-
-    expect(mockDeleteCellCallback).toHaveBeenCalledWith({
-      cellId: cellId("second-cell"),
-    });
-    expect(store.get(stagedAICellsAtom).has(cellId("first-cell"))).toBe(true);
-    expect(store.get(stagedAICellsAtom).has(cellId("second-cell"))).toBe(false);
-  });
-
-  it("should remove a provisional marimo import when a retry no longer needs it", () => {
-    const { result } = renderGenerationHook();
-    result.current.beginStagedCellGeneration();
-
-    vi.mocked(CellId.create)
-      .mockReturnValueOnce(cellId("generated-cell"))
-      .mockReturnValueOnce(cellId("marimo-import"));
-
-    result.current.onData({
-      type: "data-notebook-cells-completion",
-      data: {
-        cells: [{ language: "sql", code: "result = mo.sql('SELECT 1')" }],
-      },
-    });
-    result.current.onData({
-      type: "data-notebook-cells-completion",
-      data: {
-        cells: [{ language: "python", code: "result = 1" }],
-      },
-    });
-
-    expect(mockDeleteCellCallback).toHaveBeenCalledWith({
-      cellId: cellId("marimo-import"),
-    });
-    expect(store.get(stagedAICellsAtom).has(cellId("marimo-import"))).toBe(
-      false,
-    );
-  });
-
-  it("should discard provisional cells when generation fails", () => {
-    const { result } = renderGenerationHook();
-    result.current.beginStagedCellGeneration();
-
-    vi.mocked(CellId.create).mockReturnValue(cellId("partial-cell"));
-    result.current.onData({
-      type: "data-notebook-cells-completion",
-      data: {
-        cells: [{ language: "python", code: "partial" }],
-      },
-    });
-    result.current.finishStagedCellGeneration(false);
-
-    expect(mockDeleteCellCallback).toHaveBeenCalledWith({
-      cellId: cellId("partial-cell"),
-    });
-    expect(store.get(stagedAICellsAtom)).toEqual(new Map());
-    expect(store.get(stagedGenerationInProgressAtom)).toBe(false);
-  });
-
-  it("should not let a superseded generation mutate active cells", () => {
-    const firstGeneration = renderGenerationHook();
-    firstGeneration.result.current.beginStagedCellGeneration();
-
-    vi.mocked(CellId.create).mockReturnValueOnce(cellId("superseded-cell"));
-    firstGeneration.result.current.onData({
-      type: "data-notebook-cells-completion",
-      data: {
-        cells: [{ language: "python", code: "old" }],
-      },
-    });
-
-    const activeGeneration = renderGenerationHook();
-    activeGeneration.result.current.beginStagedCellGeneration();
-    expect(mockDeleteCellCallback).toHaveBeenCalledWith({
-      cellId: cellId("superseded-cell"),
-    });
-    mockDeleteCellCallback.mockClear();
-
-    vi.mocked(CellId.create).mockReturnValueOnce(cellId("active-cell"));
-    activeGeneration.result.current.onData({
-      type: "data-notebook-cells-completion",
-      data: {
-        cells: [{ language: "python", code: "active" }],
-      },
-    });
-
-    firstGeneration.result.current.finishStagedCellGeneration(true);
-    firstGeneration.result.current.onData({
-      type: "data-notebook-cells-completion",
-      data: {
-        cells: [{ language: "python", code: "stale" }],
-      },
-    });
-
-    expect(mockDeleteCellCallback).not.toHaveBeenCalledWith({
-      cellId: cellId("active-cell"),
-    });
-    expect(store.get(stagedAICellsAtom)).toEqual(
-      new Map([[cellId("active-cell"), { type: "add_cell" }]]),
-    );
-    expect(store.get(stagedGenerationInProgressAtom)).toBe(true);
-  });
-});

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.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.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.",
+    });
+  }
+}