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

package/src/components/editor/links/cell-link.tsx

@@ -5,6 +5,7 @@ import { useCellIds, useCellNames } from "@/core/cells/cells";
 import { useCellActions } from "@/core/cells/actions";
 import { type CellId, HTMLCellId, SCRATCH_CELL_ID } from "@/core/cells/ids";
 import { displayCellName } from "@/core/cells/names";
+import { scrollToCell } from "@/core/cells/scrollCellIntoView";
 import { goToCellLine } from "@/core/codemirror/go-to-definition/utils";
 import { useFilename } from "@/core/saving/filename";
 import { cn } from "@/utils/cn";
@@ -106,16 +107,15 @@ export function scrollAndHighlightCell(
   skipScroll?: boolean,
 ): boolean {
   const cellHtmlId = HTMLCellId.create(cellId);
-  const cell: HTMLElement | null = document.getElementById(cellHtmlId);
+  const cell = skipScroll
+    ? document.getElementById(cellHtmlId)
+    : scrollToCell(cellId);
   const isCellErrored = cell?.classList.contains("has-error");
 
   if (cell === null) {
     Logger.error(`Cell ${cellHtmlId} not found on page.`);
     return false;
   }
-  if (!skipScroll) {
-    cell.scrollIntoView({ behavior: "smooth", block: "center" });
-  }
 
   if (variant === "destructive" || (isCellErrored && variant === undefined)) {
     cell.classList.add("error-outline");

package/src/components/editor/notebook-cell.tsx

@@ -78,10 +78,7 @@ import { CollapsedCellBanner, CollapseToggle } from "./cell/collapse";
 import { DeleteButton } from "./cell/DeleteButton";
 import { PendingDeleteConfirmation } from "./cell/PendingDeleteConfirmation";
 import { RunButton } from "./cell/RunButton";
-import {
-  StagedAICellBackground,
-  StagedAICellFooter,
-} from "./cell/StagedAICell";
+import { StagedAIDeletionMarker } from "./cell/StagedAICell";
 import { useDeleteCellCallback } from "./cell/useDeleteCell";
 import { useRunCell } from "./cell/useRunCells";
 import { cellDomProps } from "./common";
@@ -674,7 +671,7 @@ const EditableCellComponent = ({
               data-has-output-above={hasOutputAbove}
               data-hidden={isMarkdownCodeHidden}
             >
-              <StagedAICellBackground cellId={cellId} />
+              <StagedAIDeletionMarker cellId={cellId} />
               <div className="absolute right-2 -top-4 z-10">
                 <CellToolbar
                   edited={cellData.edited}
@@ -837,7 +834,6 @@ const EditableCellComponent = ({
               />
             )}
             <PendingDeleteConfirmation cellId={cellId} />
-            {!isPresenting && <StagedAICellFooter cellId={cellId} />}
           </div>
           {isCollapsed && !isPresenting && (
             <CollapsedCellBanner
@@ -1198,10 +1194,7 @@ const SetupCellComponent = ({
               data-has-output-above={false}
               data-hidden={!isCellCodeShown}
             >
-              <StagedAICellBackground
-                cellId={cellId}
-                className="mo-ai-setup-cell"
-              />
+              <StagedAIDeletionMarker cellId={cellId} />
               <div className="absolute right-2 -top-4 z-10">
                 <CellToolbar
                   edited={cellData.edited}
@@ -1326,7 +1319,6 @@ const SetupCellComponent = ({
               cellId={cellId}
               debuggerActive={cellRuntime.debuggerActive}
             />
-            <StagedAICellFooter cellId={cellId} />
           </div>
         </div>
       </CellActionsContextMenu>

package/src/components/ui/toast.tsx

@@ -29,7 +29,7 @@ const VARIANT_CLASSES = {
   default:
     "bg-background border forced-colors:border-[CanvasText] forced-colors:bg-[Canvas] forced-colors:text-[CanvasText]",
   danger:
-    "group destructive text-error border-destructive bg-(--red-1) shadow-md-solid shadow-error forced-colors:border-[CanvasText] forced-colors:bg-[Canvas] forced-colors:text-[CanvasText]",
+    "group destructive text-(--red-11) border-destructive bg-(--red-1) shadow-md-solid shadow-error forced-colors:border-[CanvasText] forced-colors:bg-[Canvas] forced-colors:text-[CanvasText]",
 } as const satisfies Record<string, string>;
 
 type ToastVariant = keyof typeof VARIANT_CLASSES;

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 };

package/src/core/cells/scrollCellIntoView.ts

@@ -160,3 +160,9 @@ export function scrollCellIntoView(cellId: CellId): void {
     );
   }
 }
+
+export function scrollToCell(cellId: CellId): HTMLElement | null {
+  const cell = document.getElementById(HTMLCellId.create(cellId));
+  cell?.scrollIntoView({ behavior: "smooth", block: "center" });
+  return cell;
+}

package/src/core/codemirror/placeholder/__tests__/extensions.test.ts

@@ -145,4 +145,63 @@ describe("clickablePlaceholderExtension", () => {
 
     view.destroy();
   });
+
+  describe("caret coordinates", () => {
+    const original = Range.prototype.getClientRects;
+    let rects: Array<Pick<DOMRect, "left" | "right" | "top" | "bottom">>;
+
+    beforeEach(() => {
+      rects = [{ left: 10, right: 50, top: 2, bottom: 18 }];
+      Range.prototype.getClientRects = () => rects as unknown as DOMRectList;
+    });
+
+    afterEach(() => {
+      Range.prototype.getClientRects = original;
+    });
+
+    function renderWidget(beforeText: string) {
+      const view = new EditorView({
+        state: EditorState.create({
+          doc: "",
+          extensions: clickablePlaceholderExtension({
+            beforeText,
+            linkText: "here",
+            afterText: "",
+            onClick: vi.fn(),
+          }),
+        }),
+      });
+      const widget = view.state
+        .facet(EditorView.decorations)
+        .flatMap((deco) => (typeof deco === "function" ? [deco(view)] : [deco]))
+        .map((set) => set.iter().value?.spec.widget)
+        .find(Boolean);
+      const dom = widget.toDOM(view);
+      document.body.append(dom);
+      return { view, widget, dom };
+    }
+
+    it("collapses the caret to the start of the first line of text", () => {
+      const { view, widget, dom } = renderWidget("Type or ");
+      expect(widget.coordsAt(dom, 0, 1)).toEqual({
+        left: 10,
+        right: 10,
+        top: 2,
+        bottom: 18,
+      });
+      dom.style.direction = "rtl";
+      expect(widget.coordsAt(dom, 0, 1)).toMatchObject({ left: 50, right: 50 });
+      dom.remove();
+      view.destroy();
+    });
+
+    it("skips empty text and falls back when nothing is laid out", () => {
+      const { view, widget, dom } = renderWidget("");
+      expect(widget.coordsAt(dom, 0, 1)).toMatchObject({ left: 10 });
+      rects = [];
+      expect(widget.coordsAt(dom, 0, 1)).toBeNull();
+      dom.remove();
+      view.destroy();
+    });
+  });
 });

package/src/core/codemirror/placeholder/extensions.ts

@@ -90,6 +90,29 @@ export function clickablePlaceholderExtension(opts: {
       return wrapper;
     }
 
+    public override coordsAt(dom: HTMLElement) {
+      // A wrapped placeholder must not make the caret span every line.
+      const walker = document.createTreeWalker(dom, NodeFilter.SHOW_TEXT);
+      let text = walker.nextNode();
+      while (text && !text.textContent) {
+        text = walker.nextNode();
+      }
+      if (!text) {
+        return null;
+      }
+      const range = document.createRange();
+      range.selectNodeContents(text);
+      const rect = range.getClientRects()[0];
+      if (!rect) {
+        return null;
+      }
+      const left =
+        window.getComputedStyle(dom).direction === "rtl"
+          ? rect.right
+          : rect.left;
+      return { left, right: left, top: rect.top, bottom: rect.bottom };
+    }
+
     public override ignoreEvent(event: Event) {
       return (
         event.target instanceof Element &&

package/src/core/session/__tests__/useNotebookSession.hook.test.tsx

@@ -39,6 +39,7 @@ import { MockNotebook } from "@/__mocks__/notebook";
 import { Deferred } from "@/utils/Deferred";
 import { cellId } from "@/__tests__/branded";
 import { alertAtom, getPackageAlert } from "@/core/alerts/state";
+import { stagedAICellsAtom } from "@/core/ai/staged-cells";
 import { notebookAtom } from "@/core/cells/cells";
 import { AppConfigSchema } from "@/core/config/config-schema";
 import { ConnectionNotice } from "@/components/editor/alerts/connection-notice";
@@ -260,6 +261,39 @@ describe("useNotebookSession document model", () => {
     await expect(model.flush()).rejects.toThrow("not started");
   });
 
+  it("removes deleted cells from pending AI review while retaining surviving edits", async () => {
+    const fake = withTransport();
+    vi.mocked(useRuntimeManager).mockReturnValue(makeRuntimeManager());
+    const store = createStore();
+    const { result, unmount } = renderConnectionHook(store);
+    await started();
+    act(() => {
+      connection(fake).open();
+      fixtureMessages("notebook-events.sse").forEach(connection(fake).receive);
+    });
+    const removed = cellId("c3");
+    const retained = cellId("k7f2");
+    act(() =>
+      store.set(
+        stagedAICellsAtom,
+        new Map([
+          [removed, { type: "add_cell" }],
+          [retained, { type: "add_cell" }],
+        ]),
+      ),
+    );
+    act(() =>
+      store.set(result.current.documentModel.editAtom, (tx) =>
+        tx.deleteCell(removed),
+      ),
+    );
+    expect(store.get(notebookAtom).cellIds).not.toContain(removed);
+    expect(store.get(stagedAICellsAtom)).toEqual(
+      new Map([[retained, { type: "add_cell" }]]),
+    );
+    unmount();
+  });
+
   it("follows the host document into the notebook's cells", async () => {
     const fake = withTransport();
     vi.mocked(useRuntimeManager).mockReturnValue(makeRuntimeManager());

package/src/core/session/useNotebookSession.tsx

@@ -10,6 +10,7 @@ import {
   notebookAtom,
   useReducerActions,
 } from "@/core/cells/cells";
+import { useStagedAICellsActions } from "@/core/ai/staged-cells";
 import { followDocument } from "@/core/cells/document";
 import { AUTOCOMPLETER } from "@/core/codemirror/completion/Autocompleter";
 import { createDocumentModel } from "@/core/notebook/document";
@@ -96,6 +97,7 @@ export function useNotebookSession(opts: {
     setLastRun,
     resetCellRuntimes,
   } = useReducerActions();
+  const { reconcileStagedCells } = useStagedAICellsActions();
   const { addCellNotification } = useRunsActions();
   const setKernelState = useSetAtom(kernelStateAtom);
   const setInitialRunCompleted = useSetAtom(initialRunCompletedAtom);
@@ -441,7 +443,10 @@ export function useNotebookSession(opts: {
     followDocument({
       store,
       model: documentModel,
-      reconcile: reconcileDocument,
+      reconcile: (document) => {
+        reconcileDocument(document);
+        reconcileStagedCells(document.order);
+      },
       signal: controller.signal,
     });
     return () => {
@@ -450,7 +455,7 @@ export function useNotebookSession(opts: {
         store.set(activeDocumentModelAtom, null);
       }
     };
-  }, [documentModel, store, reconcileDocument]);
+  }, [documentModel, store, reconcileDocument, reconcileStagedCells]);
 
   // Subscribe before passive connection setup: static/WASM transports can emit synchronously.
   useLayoutEffect(() => {