@marimo-team/islands 0.25.3-dev31 → 0.25.3-dev310

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

@@ -5,24 +5,19 @@ 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 { getCellEditorView } from "../../cells/cells";
-import { updateEditorCodeFromPython } from "../../codemirror/language/utils";
+import type { CellId } from "@/core/cells/ids";
 import {
+  type Edit,
   stagedAICellsAtom,
-  stagedGenerationInProgressAtom,
   useStagedAICell,
-  useStagedCellGeneration,
   useStagedCells,
   visibleForTesting,
 } from "../staged-cells";
 
-const { stagedGenerationAtom, createActions, reducer, initialState } =
-  visibleForTesting;
+const { createActions, reducer, initialState } = visibleForTesting;
 
 // Mock the dependencies
-const mockCreateNewCell = vi.fn();
-const mockUpdateCellCode = vi.fn();
+const mockSetCellCode = vi.fn();
 const mockDeleteCellCallback = vi.fn();
 
 // Mock cell handle with editor view
@@ -34,17 +29,15 @@ const mockCellHandle = {
   },
 };
 
+vi.mock("../../cells/actions", () => ({
+  useCellActions: () => ({
+    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),
   };
 });
@@ -54,18 +47,11 @@ vi.mock("@/components/editor/cell/useDeleteCell", () => ({
 }));
 
 vi.mock("@/core/network/requests", () => ({
-  getRequestClient: () => ({ sendRun: vi.fn() }),
-}));
-
-vi.mock("../../codemirror/language/utils", () => ({
-  updateEditorCodeFromPython: vi.fn(),
+  getRequestClient: () => ({}),
 }));
 
-// Mock CellId.create
-vi.mock("@/core/cells/ids", () => ({
-  CellId: {
-    create: vi.fn(),
-  },
+vi.mock("../../codemirror/language/cell-code", () => ({
+  prepareCellCodeUpdate: vi.fn(() => ({ effects: [] })),
 }));
 
 describe("staged-cells", () => {
@@ -80,11 +66,9 @@ describe("staged-cells", () => {
 
     // Reset mocks
     vi.clearAllMocks();
-    vi.mocked(CellId.create).mockReset();
 
     // Reset the atom state
     store.set(stagedAICellsAtom, new Map());
-    store.set(stagedGenerationAtom, null);
   });
 
   describe("reducer and actions", () => {
@@ -150,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 }],
@@ -249,594 +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(mockCreateNewCell).toHaveBeenCalledWith({
-        cellId: "__end__",
-        code: testCode,
-        before: false,
-        newCellId: mockCellId,
-      });
-    });
-
-    it("should delete a staged cell", () => {
-      const { result } = renderHook(() => useStagedCells(store));
-      const testCellId = cellId("test-cell-id");
-
-      result.current.deleteStagedCell(testCellId);
-
-      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({
-        type: "update_cell",
-        previousCode: "old code",
-      });
-    });
-
-    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("should handle multiple operations correctly", () => {
-      const { result } = renderHook(() => useStagedCells(store));
+  describe("useStagedCells", () => {
+    const stage = (id: CellId, edit: Edit) =>
+      store.set(stagedAICellsAtom, new Map([[id, edit]]));
 
-      // 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(mockCreateNewCell).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("should track edit history for updated cells", () => {
-      const { result } = renderHook(() => useStagedCells(store));
-
-      // Add a cell with update_cell edit type
-      result.current.addStagedCell({
-        cellId: cellId1,
-        edit: { type: "update_cell", previousCode: "previous code" },
-      });
-
-      const state = store.get(stagedAICellsAtom);
-      const edit = state.get(cellId1);
-      expect(edit).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: "previous code",
+        previousCode: "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({
+      expect(store.get(stagedAICellsAtom).size).toBe(0);
+      expect(mockDeleteCellCallback).not.toHaveBeenCalled();
+    });
+
+    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("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("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: "delete_cell", previousCode: "deleted content" },
+        code: "x = 1",
       });
-
-      const state = store.get(stagedAICellsAtom);
-      const edit = state.get(cellId1);
-      expect(edit).toEqual({
-        type: "delete_cell",
-        previousCode: "deleted content",
-      });
-    });
-  });
-});
-
-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(mockCreateNewCell).not.toHaveBeenCalled();
-    expect(updateEditorCodeFromPython).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(mockCreateNewCell).toHaveBeenCalledWith({
-      cellId: "__end__",
-      code: "print('```')",
-      before: false,
-      newCellId: "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(mockCreateNewCell).toHaveBeenCalledWith(
-      expect.objectContaining({
-        code: "value = 1",
-        newCellId: "first-cell",
-      }),
-    );
-    expect(mockCreateNewCell).toHaveBeenCalledWith(
-      expect.objectContaining({
-        code: "value",
-        newCellId: "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(mockCreateNewCell).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);
-  });
-
-  it("should update notebook state 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(mockUpdateCellCode).toHaveBeenCalledWith({
-      cellId: cellId("unmounted-cell"),
-      code: "value = 1",
-      formattingChange: false,
-    });
-  });
-
-  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" }],
-      },
-    });
+    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 });
 
-    expect(mockDeleteCellCallback).not.toHaveBeenCalledWith({
-      cellId: cellId("active-cell"),
+      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);
     });
-    expect(store.get(stagedAICellsAtom)).toEqual(
-      new Map([[cellId("active-cell"), { type: "add_cell" }]]),
-    );
-    expect(store.get(stagedGenerationInProgressAtom)).toBe(true);
   });
 });