@marimo-team/islands 0.25.3-dev266 → 0.25.3-dev268

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

@@ -5,8 +5,9 @@ import { getDefaultStore, Provider } from "jotai";
 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 type { CellId } from "@/core/cells/ids";
 import {
+  type Edit,
   stagedAICellsAtom,
   useStagedAICell,
   useStagedCells,
@@ -16,7 +17,6 @@ import {
 const { createActions, reducer, initialState } = visibleForTesting;
 
 // Mock the dependencies
-const mockCreateCell = vi.fn();
 const mockSetCellCode = vi.fn();
 const mockDeleteCellCallback = vi.fn();
 
@@ -31,7 +31,6 @@ const mockCellHandle = {
 
 vi.mock("../../cells/actions", () => ({
   useCellActions: () => ({
-    createCell: mockCreateCell,
     setCellCode: mockSetCellCode,
   }),
 }));
@@ -55,13 +54,6 @@ vi.mock("../../codemirror/language/cell-code", () => ({
   prepareCellCodeUpdate: vi.fn(() => ({ effects: [] })),
 }));
 
-// Mock CellId.create
-vi.mock("@/core/cells/ids", () => ({
-  CellId: {
-    create: vi.fn(),
-  },
-}));
-
 describe("staged-cells", () => {
   let store: ReturnType<typeof getDefaultStore>;
   let cellId1: CellId;
@@ -74,7 +66,6 @@ describe("staged-cells", () => {
 
     // Reset mocks
     vi.clearAllMocks();
-    vi.mocked(CellId.create).mockReset();
 
     // Reset the atom state
     store.set(stagedAICellsAtom, new Map());
@@ -143,6 +134,31 @@ describe("staged-cells", () => {
       });
     });
 
+    it("drops a replaced cell while preserving the four remaining pending edits", () => {
+      const original = cellId("original");
+      const remaining = [cellId("a"), cellId("b"), cellId("c"), cellId("d")];
+      const state = new Map<CellId, Edit>([
+        [original, { type: "add_cell" }],
+        [remaining[0], { type: "add_cell" }],
+        [remaining[1], { type: "add_cell" }],
+        [remaining[2], { type: "update_cell", previousCode: "x = 1" }],
+        [remaining[3], { type: "delete_cell", previousCode: "y = 2" }],
+      ]);
+      const next = reducer(state, {
+        type: "reconcileStagedCells",
+        payload: remaining,
+      });
+      expect(next).toEqual(
+        new Map([...state].filter(([id]) => id !== original)),
+      );
+      expect(
+        reducer(next, { type: "reconcileStagedCells", payload: remaining }),
+      ).toBe(next);
+      expect(
+        reducer(next, { type: "reconcileStagedCells", payload: [] }),
+      ).toEqual(new Map());
+    });
+
     it("should remove cell IDs", () => {
       const state = new Map([
         [cellId1, { type: "add_cell" as const }],
@@ -242,134 +258,66 @@ describe("staged-cells", () => {
     });
   });
 
-  describe("useStagedCells hook", () => {
-    it("should create a staged cell with code", () => {
-      const { result } = renderHook(() => useStagedCells(store));
-      const testCode = "print('hello world')";
-
-      // Mock CellId.create to return a predictable ID
-      const mockCellId = cellId("mock-cell-id");
-      vi.mocked(CellId.create).mockReturnValue(mockCellId);
-
-      const returnedCellId = result.current.createStagedCell(testCode);
-
-      expect(returnedCellId).toBe(mockCellId);
-      expect(mockCreateCell).toHaveBeenCalledWith({
-        code: testCode,
-        id: mockCellId,
-      });
-    });
-
-    it("should delete a staged cell", () => {
-      const { result } = renderHook(() => useStagedCells(store));
-      const testCellId = cellId("test-cell-id");
-
-      result.current.deleteStagedCell(testCellId);
+  describe("useStagedCells", () => {
+    const stage = (id: CellId, edit: Edit) =>
+      store.set(stagedAICellsAtom, new Map([[id, edit]]));
 
-      expect(mockDeleteCellCallback).toHaveBeenCalledWith({
-        cellId: testCellId,
-      });
-    });
-
-    it("should add staged cell with edit info", () => {
-      const { result } = renderHook(() => useStagedCells(store));
-
-      result.current.addStagedCell({
-        cellId: cellId1,
-        edit: { type: "update_cell", previousCode: "old code" },
-      });
-
-      // Check that the cell was added to the atom with edit info
-      const state = store.get(stagedAICellsAtom);
-      expect(state.has(cellId1)).toBe(true);
-      expect(state.get(cellId1)).toEqual({
+    it("keeps added and updated cells by unstaging them", () => {
+      const { result } = renderHook(() => useStagedCells());
+      stage(cellId1, { type: "update_cell", previousCode: "x = 1" });
+      result.current.acceptStagedCell(cellId1, {
         type: "update_cell",
-        previousCode: "old code",
+        previousCode: "x = 1",
       });
+      expect(store.get(stagedAICellsAtom).size).toBe(0);
+      expect(mockDeleteCellCallback).not.toHaveBeenCalled();
     });
 
-    it("should remove staged cell", () => {
-      const { result } = renderHook(() => useStagedCells(store));
-
-      // First add cells
-      result.current.addStagedCell({
-        cellId: cellId1,
-        edit: { type: "add_cell" },
-      });
-      result.current.addStagedCell({
-        cellId: cellId2,
-        edit: { type: "add_cell" },
-      });
-
-      // Then remove one
-      result.current.removeStagedCell(cellId1);
-
-      // Check that only the remaining cell is in the map
-      const state = store.get(stagedAICellsAtom);
-      expect(state.has(cellId1)).toBe(false);
-      expect(state.has(cellId2)).toBe(true);
+    it("deletes the cell when a proposed deletion is kept", () => {
+      const { result } = renderHook(() => useStagedCells());
+      const edit: Edit = { type: "delete_cell", previousCode: "x = 1" };
+      stage(cellId1, edit);
+      result.current.acceptStagedCell(cellId1, edit);
+      expect(mockDeleteCellCallback).toHaveBeenCalledWith({ cellId: cellId1 });
+      expect(store.get(stagedAICellsAtom).size).toBe(0);
     });
 
-    it("should handle multiple operations correctly", () => {
-      const { result } = renderHook(() => useStagedCells(store));
-
-      // Create a staged cell
-      const mockCellId = cellId("mock-cell-id");
-      vi.mocked(CellId.create).mockReturnValue(mockCellId);
-
-      const createdCellId = result.current.createStagedCell("test code");
-
-      // Verify it was created and added
-      expect(createdCellId).toBe(mockCellId);
-      expect(mockCreateCell).toHaveBeenCalled();
-
-      let state = store.get(stagedAICellsAtom);
-      expect(state.has(mockCellId)).toBe(true);
-      expect(state.get(mockCellId)).toEqual({ type: cellId("add_cell") });
-
-      // Delete the staged cell
-      result.current.deleteStagedCell(mockCellId);
-      expect(mockDeleteCellCallback).toHaveBeenCalledWith({
-        cellId: mockCellId,
-      });
-
-      // Verify it was removed from staged cells
-      state = store.get(stagedAICellsAtom);
-      expect(state.has(mockCellId)).toBe(false);
+    it("restores the previous code when an update is discarded", () => {
+      const { result } = renderHook(() => useStagedCells());
+      const edit: Edit = { type: "update_cell", previousCode: "x = 1" };
+      stage(cellId1, edit);
+      result.current.rejectStagedCell(cellId1, edit);
+      expect(
+        mockCellHandle.current.editorViewOrNull.dispatch,
+      ).toHaveBeenCalled();
+      expect(store.get(stagedAICellsAtom).size).toBe(0);
     });
 
-    it("should track edit history for updated cells", () => {
-      const { result } = renderHook(() => useStagedCells(store));
-
-      // Add a cell with update_cell edit type
-      result.current.addStagedCell({
+    it("writes to the document when the editor is not mounted", async () => {
+      const { getCellEditorView } = await import("../../cells/cells");
+      vi.mocked(getCellEditorView).mockReturnValueOnce(undefined);
+      const { result } = renderHook(() => useStagedCells());
+      const edit: Edit = { type: "update_cell", previousCode: "x = 1" };
+      stage(cellId1, edit);
+      result.current.rejectStagedCell(cellId1, edit);
+      expect(mockSetCellCode).toHaveBeenCalledWith({
         cellId: cellId1,
-        edit: { type: "update_cell", previousCode: "previous code" },
-      });
-
-      const state = store.get(stagedAICellsAtom);
-      const edit = state.get(cellId1);
-      expect(edit).toEqual({
-        type: "update_cell",
-        previousCode: "previous code",
+        code: "x = 1",
       });
     });
 
-    it("should track edit history for deleted cells", () => {
-      const { result } = renderHook(() => useStagedCells(store));
-
-      // Add a cell with delete_cell edit type
-      result.current.addStagedCell({
-        cellId: cellId1,
-        edit: { type: "delete_cell", previousCode: "deleted content" },
-      });
-
-      const state = store.get(stagedAICellsAtom);
-      const edit = state.get(cellId1);
-      expect(edit).toEqual({
-        type: "delete_cell",
-        previousCode: "deleted content",
-      });
+    it("deletes an added cell and leaves a proposed deletion in place when discarded", () => {
+      const { result } = renderHook(() => useStagedCells());
+      stage(cellId1, { type: "add_cell" });
+      result.current.rejectStagedCell(cellId1, { type: "add_cell" });
+      expect(mockDeleteCellCallback).toHaveBeenCalledWith({ cellId: cellId1 });
+
+      mockDeleteCellCallback.mockClear();
+      const deletion: Edit = { type: "delete_cell", previousCode: "x = 1" };
+      stage(cellId2, deletion);
+      result.current.rejectStagedCell(cellId2, deletion);
+      expect(mockDeleteCellCallback).not.toHaveBeenCalled();
+      expect(store.get(stagedAICellsAtom).size).toBe(0);
     });
   });
 });

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

@@ -4,13 +4,11 @@ import { useAtomValue } from "jotai";
 import { selectAtom } from "jotai/utils";
 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 { 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 chat and agent workflows. */
@@ -36,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;
@@ -48,44 +58,27 @@ 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 { setCellCode, createCell } = 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" } });
-    createCell({ code, id: 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;
-    }
-
-    // The AI may change the cell's language, and only the editor infers
-    // one. Without a mounted editor the code arrives through the document.
+  // 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) {
       const update = prepareCellCodeUpdate(editorView.state, code, {
@@ -104,13 +97,30 @@ export function useStagedCells(store: JotaiStore) {
     deleteCellCallback({ cellId });
   };
 
-  return {
-    createStagedCell,
-    updateStagedCell,
-    addStagedCell,
-    removeStagedCell,
-    deleteStagedCell,
+  const acceptStagedCell = (cellId: CellId, edit: Edit) => {
+    if (edit.type === "delete_cell") {
+      deleteStagedCell(cellId);
+    } else {
+      removeStagedCell(cellId);
+    }
+  };
+
+  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;
+    }
   };
+
+  return { acceptStagedCell, rejectStagedCell };
 }
 
 export { stagedAICellsAtom };