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

package/src/components/editor/chrome/wrapper/__tests__/pending-ai-cells.test.tsx

@@ -4,16 +4,23 @@ import { fireEvent, render, screen } from "@testing-library/react";
 import { createStore, Provider } from "jotai";
 import { describe, expect, it, vi } from "vitest";
 import type { ChatStatus } from "ai";
-import { useChatGeneration } from "@/components/chat/use-chat-generation";
+import {
+  isChatBusy,
+  useChatGeneration,
+} from "@/components/chat/use-chat-generation";
 import { cellId } from "@/__tests__/branded";
+import { HTMLCellId } from "@/core/cells/ids";
+import { scrollToCell } from "@/core/cells/scrollCellIntoView";
 import { TooltipProvider } from "@/components/ui/tooltip";
 import { stagedAICellsAtom } from "@/core/ai/staged-cells";
-import { scrollAndHighlightCell } from "../../../links/cell-link";
-import { StagedAICellFooter } from "../../../cell/StagedAICell";
 import { PendingAICells } from "../pending-ai-cells";
+import { PendingAIChangesBadge } from "../pending-ai-changes-badge";
+
+const deleteCell = vi.hoisted(() => vi.fn());
+const setCellCode = vi.hoisted(() => vi.fn());
 
 vi.mock("@/components/editor/cell/useDeleteCell", () => ({
-  useDeleteCellCallback: () => vi.fn(),
+  useDeleteCellCallback: () => deleteCell,
 }));
 
 vi.mock("@/core/cells/cells", () => ({
@@ -22,7 +29,7 @@ vi.mock("@/core/cells/cells", () => ({
 vi.mock("@/core/cells/actions", () => ({
   useCellActions: () => ({
     createCell: vi.fn(),
-    setCellCode: vi.fn(),
+    setCellCode,
   }),
 }));
 
@@ -31,8 +38,8 @@ vi.mock("@/components/editor/cell/useRunCells", () => ({
   useRunCell: () => vi.fn(),
 }));
 
-vi.mock("@/components/editor/links/cell-link", () => ({
-  scrollAndHighlightCell: vi.fn(),
+vi.mock("@/core/cells/scrollCellIntoView", () => ({
+  scrollToCell: vi.fn(),
 }));
 
 describe("PendingAICells", () => {
@@ -63,60 +70,60 @@ describe("PendingAICells", () => {
       name: "Previous pending AI cell",
     });
 
-    fireEvent.click(next);
-    expect(scrollAndHighlightCell).toHaveBeenLastCalledWith(first, "focus");
-    expect(screen.getByText("1 / 3")).toBeInTheDocument();
+    fireEvent.click(screen.getByRole("button", { name: "Cell 1 of 3" }));
+    expect(scrollToCell).toHaveBeenLastCalledWith(first);
+    expect(screen.getByText("Cell 1 of 3")).toBeInTheDocument();
 
     fireEvent.click(next);
-    expect(scrollAndHighlightCell).toHaveBeenLastCalledWith(second, "focus");
-    expect(screen.getByText("2 / 3")).toBeInTheDocument();
+    expect(scrollToCell).toHaveBeenLastCalledWith(second);
+    expect(screen.getByText("Cell 2 of 3")).toBeInTheDocument();
 
     fireEvent.click(previous);
-    expect(scrollAndHighlightCell).toHaveBeenLastCalledWith(first, "focus");
-    expect(screen.getByText("1 / 3")).toBeInTheDocument();
+    expect(scrollToCell).toHaveBeenLastCalledWith(first);
+    expect(screen.getByText("Cell 1 of 3")).toBeInTheDocument();
 
     fireEvent.click(next);
-    expect(scrollAndHighlightCell).toHaveBeenLastCalledWith(second, "focus");
-    expect(screen.getByText("2 / 3")).toBeInTheDocument();
+    expect(scrollToCell).toHaveBeenLastCalledWith(second);
+    expect(screen.getByText("Cell 2 of 3")).toBeInTheDocument();
 
     fireEvent.click(next);
-    expect(scrollAndHighlightCell).toHaveBeenLastCalledWith(third, "focus");
-    expect(screen.getByText("3 / 3")).toBeInTheDocument();
+    expect(scrollToCell).toHaveBeenLastCalledWith(third);
+    expect(screen.getByText("Cell 3 of 3")).toBeInTheDocument();
 
     fireEvent.click(next);
-    expect(scrollAndHighlightCell).toHaveBeenLastCalledWith(first, "focus");
-    expect(screen.getByText("1 / 3")).toBeInTheDocument();
+    expect(scrollToCell).toHaveBeenLastCalledWith(first);
+    expect(screen.getByText("Cell 1 of 3")).toBeInTheDocument();
 
     fireEvent.click(previous);
-    expect(scrollAndHighlightCell).toHaveBeenLastCalledWith(third, "focus");
-    expect(screen.getByText("3 / 3")).toBeInTheDocument();
+    expect(scrollToCell).toHaveBeenLastCalledWith(third);
+    expect(screen.getByText("Cell 3 of 3")).toBeInTheDocument();
   });
 });
 
 function ChatGeneration({ status }: { status: ChatStatus }) {
-  useChatGeneration(status);
+  useChatGeneration(isChatBusy(status));
   return null;
 }
 
 it("locks staged edits during generation and unlocks after completion, stop, error, or unmount", () => {
   const store = createStore();
   const pendingCell = cellId("pending-cell");
-  store.set(stagedAICellsAtom, new Map([[pendingCell, { type: "add_cell" }]]));
+  store.set(
+    stagedAICellsAtom,
+    new Map([
+      [pendingCell, { type: "add_cell" }],
+      [cellId("other-pending-cell"), { type: "add_cell" }],
+    ]),
+  );
   const ui = (status?: ChatStatus) => (
     <Provider store={store}>
       <TooltipProvider>
         {status && <ChatGeneration status={status} />}
         <PendingAICells />
-        <StagedAICellFooter cellId={pendingCell} />
       </TooltipProvider>
     </Provider>
   );
   const { rerender } = render(ui("ready"));
-  const keep = screen.getByRole("button", { name: "Keep all" });
-  const discard = screen.getByRole("button", { name: "Discard all" });
-  const keepCell = screen.getByRole("button", { name: "Keep cell" });
-  const discardCell = screen.getByRole("button", { name: "Discard cell" });
-  const keepAndRun = keep.parentElement!.querySelectorAll("button")[1];
   for (const status of [
     "submitted",
     "streaming",
@@ -130,15 +137,169 @@ it("locks staged edits during generation and unlocks after completion, stop, err
   ] as const) {
     rerender(ui(status));
     const isGenerating = status === "submitted" || status === "streaming";
-    for (const button of [keep, discard, keepAndRun, keepCell, discardCell]) {
+    const keep = screen.getByRole("button", { name: "Keep all" });
+    expect(keep.hasAttribute("disabled")).toBe(isGenerating);
+    expect(
+      screen.queryByRole("button", { name: "Keep all" }),
+    ).toBeInTheDocument();
+    const keepCell = screen.getByRole("button", { name: "Keep" });
+    const discardCell = screen.getByRole("button", { name: "Discard" });
+    for (const button of [keepCell, discardCell]) {
       expect(button.hasAttribute("disabled")).toBe(isGenerating);
       if (isGenerating) {
         fireEvent.click(button);
       }
     }
     expect(store.get(stagedAICellsAtom).has(pendingCell)).toBe(true);
+    expect(screen.queryByRole("button", { name: "Keep" })).toBeInTheDocument();
   }
 
-  fireEvent.click(keep);
+  fireEvent.click(screen.getByRole("button", { name: "Keep all" }));
   expect(store.get(stagedAICellsAtom).size).toBe(0);
 });
+
+it.each(["Delete cell", "Keep cell"])(
+  "resolves a proposed deletion with %s and advances review",
+  (label) => {
+    deleteCell.mockClear();
+    const store = createStore();
+    const first = cellId("deleted-cell");
+    const second = cellId("added-cell");
+    store.set(
+      stagedAICellsAtom,
+      new Map([
+        [first, { type: "delete_cell", previousCode: "x = 1" }],
+        [second, { type: "add_cell" }],
+      ]),
+    );
+    render(
+      <Provider store={store}>
+        <PendingAICells />
+      </Provider>,
+    );
+    fireEvent.click(screen.getByRole("button", { name: label }));
+    expect(store.get(stagedAICellsAtom).has(first)).toBe(false);
+    expect(deleteCell).toHaveBeenCalledTimes(label === "Delete cell" ? 1 : 0);
+    expect(scrollToCell).toHaveBeenLastCalledWith(second);
+    expect(screen.getByText("1 cell changed")).toBeInTheDocument();
+    expect(
+      screen.queryByRole("button", { name: "Keep all" }),
+    ).not.toBeInTheDocument();
+    expect(
+      screen.queryByRole("button", { name: "Discard all" }),
+    ).not.toBeInTheDocument();
+    fireEvent.click(screen.getByRole("button", { name: "Keep" }));
+    expect(
+      screen.queryByRole("region", { name: "Review AI changes" }),
+    ).not.toBeInTheDocument();
+  },
+);
+
+it("keeps batch discard visible and locks it during generation", () => {
+  deleteCell.mockClear();
+  const store = createStore();
+  const pendingCell = cellId("discarded-cell");
+  store.set(
+    stagedAICellsAtom,
+    new Map([
+      [pendingCell, { type: "add_cell" }],
+      [cellId("other-pending-cell"), { type: "add_cell" }],
+    ]),
+  );
+  const ui = (status: ChatStatus) => (
+    <Provider store={store}>
+      <ChatGeneration status={status} />
+      <PendingAICells />
+    </Provider>
+  );
+  const { rerender } = render(ui("streaming"));
+  const discard = screen.getByRole("button", { name: "Discard all" });
+  expect(discard).toBeVisible();
+  expect(discard).toBeDisabled();
+  fireEvent.click(discard);
+  expect(store.get(stagedAICellsAtom).has(pendingCell)).toBe(true);
+  rerender(ui("ready"));
+  expect(discard).toBeEnabled();
+  fireEvent.click(discard);
+  expect(deleteCell).toHaveBeenCalledTimes(2);
+  expect(deleteCell).toHaveBeenCalledWith({ cellId: pendingCell });
+  expect(store.get(stagedAICellsAtom).size).toBe(0);
+});
+
+it("clears the panel badge when all pending changes are accepted", () => {
+  const store = createStore();
+  store.set(
+    stagedAICellsAtom,
+    new Map([[cellId("pending"), { type: "add_cell" }]]),
+  );
+  render(
+    <Provider store={store}>
+      <PendingAIChangesBadge />
+      <PendingAICells />
+    </Provider>,
+  );
+  expect(screen.getByText("1 pending AI change")).toBeInTheDocument();
+  fireEvent.click(screen.getByRole("button", { name: "Keep" }));
+  expect(screen.queryByText("1 pending AI change")).not.toBeInTheDocument();
+});
+
+it("discards an edit to a cell whose editor is not mounted", () => {
+  setCellCode.mockClear();
+  const store = createStore();
+  const collapsed = cellId("collapsed-cell");
+  store.set(
+    stagedAICellsAtom,
+    new Map([[collapsed, { type: "update_cell", previousCode: "x = 1" }]]),
+  );
+  render(
+    <Provider store={store}>
+      <PendingAICells />
+    </Provider>,
+  );
+  fireEvent.click(screen.getByRole("button", { name: "Discard" }));
+  expect(setCellCode).toHaveBeenCalledWith({
+    cellId: collapsed,
+    code: "x = 1",
+  });
+  expect(store.get(stagedAICellsAtom).size).toBe(0);
+});
+
+it("outlines only the reviewed cell, moving the outline with the selection", () => {
+  const store = createStore();
+  const first = cellId("outlined-first");
+  const second = cellId("outlined-second");
+  store.set(
+    stagedAICellsAtom,
+    new Map([
+      [first, { type: "add_cell" }],
+      [second, { type: "add_cell" }],
+    ]),
+  );
+  const cells = [first, second].map((id) => {
+    const element = document.createElement("div");
+    element.id = HTMLCellId.create(id);
+    document.body.append(element);
+    return element;
+  });
+  const { unmount } = render(
+    <Provider store={store}>
+      <TooltipProvider>
+        <PendingAICells />
+      </TooltipProvider>
+    </Provider>,
+  );
+  const outlined = () =>
+    cells.map((cell) => cell.hasAttribute("data-ai-review-selected"));
+
+  expect(outlined()).toEqual([false, false]);
+  fireEvent.click(screen.getByRole("button", { name: "Cell 1 of 2" }));
+  expect(outlined()).toEqual([true, false]);
+  fireEvent.click(screen.getByRole("button", { name: "Next pending AI cell" }));
+  expect(outlined()).toEqual([false, true]);
+
+  unmount();
+  expect(outlined()).toEqual([false, false]);
+  for (const cell of cells) {
+    cell.remove();
+  }
+});

package/src/components/editor/chrome/wrapper/app-chrome.tsx

@@ -68,7 +68,7 @@ import {
 import { BackendConnectionStatus } from "./footer-items/backend-status";
 import { LspStatus } from "./footer-items/lsp-status";
 import { PanelsWrapper } from "./panels";
-import { PendingAICells } from "./pending-ai-cells";
+import { PendingAIChangesBadge } from "./pending-ai-changes-badge";
 import { useAiPanelTab } from "./useAiPanel";
 import { useDependencyPanelTab } from "./useDependencyPanelTab";
 import { handleDragging } from "./utils";
@@ -558,6 +558,7 @@ export const AppChrome: React.FC<PropsWithChildren> = ({ children }) => {
                   )}
                 />
                 {panel.label}
+                {panel.type === "ai" && <PendingAIChangesBadge />}
               </div>
             )}
           />
@@ -618,7 +619,6 @@ export const AppChrome: React.FC<PropsWithChildren> = ({ children }) => {
         </Panel>
         <ContextAwarePanel />
       </PanelGroup>
-      <PendingAICells />
       <ErrorBoundary>
         <TooltipProvider>
           <Footer />

package/src/components/editor/chrome/wrapper/pending-ai-cells.tsx

@@ -1,38 +1,47 @@
 /* Copyright 2026 Marimo. All rights reserved. */
 
 import { useAtomValue, useStore } from "jotai";
-import { ChevronDown, ChevronUp, SparklesIcon } from "lucide-react";
-import { useState } from "react";
+import { ChevronLeft, ChevronRight } from "lucide-react";
+import { useEffect, useState } from "react";
 import { Button } from "@/components/ui/button";
-import { stagedAICellsAtom, useStagedCells } from "@/core/ai/staged-cells";
+import {
+  type Edit,
+  stagedAICellsAtom,
+  useStagedCells,
+} from "@/core/ai/staged-cells";
 import { chatGenerationInProgressAtom } from "@/core/ai/state";
-import type { CellId } from "@/core/cells/ids";
+import { type CellId, HTMLCellId } from "@/core/cells/ids";
+import { scrollToCell } from "@/core/cells/scrollCellIntoView";
 import { getNextIndex } from "@/utils/arrays";
-import { cn } from "@/utils/cn";
-import {
-  AcceptCompletionButton,
-  RejectCompletionButton,
-} from "../../ai/completion-handlers";
-import { acceptStagedCell, rejectStagedCell } from "../../cell/StagedAICell";
-import { useRunCells } from "../../cell/useRunCells";
-import { scrollAndHighlightCell } from "../../links/cell-link";
+
+const REVIEW_LABELS: Record<Edit["type"], { accept: string; reject: string }> =
+  {
+    add_cell: { accept: "Keep", reject: "Discard" },
+    update_cell: { accept: "Keep", reject: "Discard" },
+    delete_cell: { accept: "Delete cell", reject: "Keep cell" },
+  };
 
 export const PendingAICells: React.FC = () => {
-  const [currentIndex, setCurrentIndex] = useState<number | null>(null);
+  const [selectedCellId, setSelectedCellId] = useState<CellId | null>(null);
 
   const stagedAiCells = useAtomValue(stagedAICellsAtom);
   const isGenerating = useAtomValue(chatGenerationInProgressAtom);
   const listStagedCells = [...stagedAiCells.keys()];
+  const selectedIndex =
+    selectedCellId === null ? -1 : listStagedCells.indexOf(selectedCellId);
+  const currentIndex = Math.max(0, selectedIndex);
   const store = useStore();
-  const { deleteStagedCell, removeStagedCell } = useStagedCells(store);
-  const runCell = useRunCells();
+  const { acceptStagedCell, rejectStagedCell } = useStagedCells();
+  const reviewedCellId = selectedIndex === -1 ? null : selectedCellId;
+  useReviewedCellMarker(reviewedCellId);
 
   if (stagedAiCells.size === 0) {
     return null;
   }
 
-  const scrollToCell = (cellId: CellId) => {
-    scrollAndHighlightCell(cellId, "focus");
+  const selectCell = (cellId: CellId) => {
+    setSelectedCellId(cellId);
+    scrollToCell(cellId);
   };
 
   const clickNext = (direction: "up" | "down") => {
@@ -41,86 +50,141 @@ export const PendingAICells: React.FC = () => {
       listStagedCells.length,
       direction,
     );
-    setCurrentIndex(newIndex);
-    scrollToCell(listStagedCells[newIndex]);
+    selectCell(listStagedCells[newIndex]);
   };
 
   const acceptAllCompletions = () => {
     for (const [cellId, edit] of stagedAiCells) {
-      acceptStagedCell(cellId, edit, removeStagedCell, deleteStagedCell);
+      acceptStagedCell(cellId, edit);
     }
   };
 
   const rejectAllCompletions = () => {
     for (const [cellId, edit] of stagedAiCells) {
-      rejectStagedCell(cellId, edit, removeStagedCell, deleteStagedCell);
+      rejectStagedCell(cellId, edit);
     }
   };
 
-  const runAllCells = () => {
-    const cellsToRun = listStagedCells.filter(
-      (cellId) => stagedAiCells.get(cellId)?.type !== "delete_cell",
-    );
-    runCell(cellsToRun);
-  };
+  const selectedCell = listStagedCells[currentIndex];
+  const selectedEdit = stagedAiCells.get(selectedCell);
+  const hasMultipleCells = listStagedCells.length > 1;
 
-  const cyanShadow = "shadow-[0_0_6px_0_#00A2C733]";
+  const resolveSelected = (decision: "accept" | "reject") => {
+    if (!selectedEdit || store.get(chatGenerationInProgressAtom)) {
+      return;
+    }
+    if (decision === "accept") {
+      acceptStagedCell(selectedCell, selectedEdit);
+    } else {
+      rejectStagedCell(selectedCell, selectedEdit);
+    }
+    if (hasMultipleCells) {
+      clickNext("up");
+    } else {
+      setSelectedCellId(null);
+    }
+  };
 
   return (
-    <div
-      className={cn(
-        "fixed bottom-16 left-1/2 transform -translate-x-1/2 z-50 bg-background/95 backdrop-blur-sm supports-backdrop-filter:bg-background/80 border border-border rounded-lg px-3 py-2 flex items-center justify-between gap-2.5 w-100",
-        cyanShadow,
-      )}
+    <section
+      aria-label="Review AI changes"
+      className="@container/review shrink-0 border-t border-border bg-background px-3 py-2 text-sm"
     >
-      <SparklesIcon className="h-4 w-4 text-primary" />
-
-      <div className="flex items-center">
-        <Button
-          aria-label="Previous pending AI cell"
-          variant="ghost"
-          size="icon"
-          onClick={() => clickNext("down")}
-        >
-          <ChevronUp className="h-3.5 w-3.5" />
-        </Button>
-        <span className="text-xs font-mono min-w-14 text-center">
-          {currentIndex === null
-            ? `${listStagedCells.length} pending`
-            : `${currentIndex + 1} / ${listStagedCells.length}`}
-        </span>
-        <Button
-          aria-label="Next pending AI cell"
-          variant="ghost"
-          size="icon"
-          onClick={() => clickNext("up")}
-        >
-          <ChevronDown className="h-3.5 w-3.5" />
-        </Button>
-      </div>
-
-      <div className="h-5 w-px bg-border" />
-
-      <div className="flex items-center gap-1.5">
-        <AcceptCompletionButton
-          multipleCompletions={true}
-          onAccept={acceptAllCompletions}
-          isLoading={isGenerating}
-          label="Keep all"
-          size="xs"
-          buttonStyles="h-6.5"
-          playButtonStyles="h-6.5"
-          runCell={runAllCells}
-        />
-        <RejectCompletionButton
-          multipleCompletions={true}
-          onDecline={rejectAllCompletions}
-          disabled={isGenerating}
-          size="xs"
-          label="Discard all"
-          className="h-6.5"
-        />
+      <div className="flex flex-wrap items-center gap-x-3 gap-y-2">
+        {hasMultipleCells && (
+          <div className="flex flex-wrap items-center gap-2">
+            <span className="whitespace-nowrap text-muted-foreground">
+              {listStagedCells.length} cells changed
+            </span>
+            <Button
+              variant="outline"
+              size="sm"
+              className="chat-control chat-selector h-7 px-2 rounded-sm font-normal"
+              onClick={acceptAllCompletions}
+              disabled={isGenerating}
+            >
+              Keep all
+            </Button>
+            <Button
+              variant="outline"
+              size="sm"
+              className="chat-control chat-selector h-7 px-2 rounded-sm font-normal"
+              onClick={rejectAllCompletions}
+              disabled={isGenerating}
+            >
+              Discard all
+            </Button>
+          </div>
+        )}
+        <div className="flex flex-wrap items-center gap-2">
+          {hasMultipleCells && (
+            <span
+              aria-hidden="true"
+              className="hidden h-4 w-px bg-border @min-[600px]/review:block"
+            />
+          )}
+          <div className="flex h-7 w-fit items-center">
+            {hasMultipleCells && (
+              <Button
+                aria-label="Previous pending AI cell"
+                variant="ghost"
+                size="icon"
+                className="chat-control size-6 rounded-sm"
+                onClick={() => clickNext("down")}
+              >
+                <ChevronLeft className="size-4" />
+              </Button>
+            )}
+            <Button
+              variant="ghost"
+              size="sm"
+              className="chat-control h-6 px-1 rounded-sm font-normal tabular-nums whitespace-nowrap"
+              onClick={() => selectCell(selectedCell)}
+            >
+              <span aria-live="polite">
+                {hasMultipleCells
+                  ? `Cell ${currentIndex + 1} of ${listStagedCells.length}`
+                  : "1 cell changed"}
+              </span>
+            </Button>
+            {hasMultipleCells && (
+              <Button
+                aria-label="Next pending AI cell"
+                variant="ghost"
+                size="icon"
+                className="chat-control size-6 rounded-sm"
+                onClick={() => clickNext("up")}
+              >
+                <ChevronRight className="size-4" />
+              </Button>
+            )}
+          </div>
+          {selectedEdit &&
+            (["accept", "reject"] as const).map((decision) => (
+              <Button
+                key={decision}
+                variant="outline"
+                size="sm"
+                className="chat-control chat-selector h-7 px-2 rounded-sm font-normal"
+                disabled={isGenerating}
+                onClick={() => resolveSelected(decision)}
+              >
+                {REVIEW_LABELS[selectedEdit.type][decision]}
+              </Button>
+            ))}
+        </div>
       </div>
-    </div>
+    </section>
   );
 };
+
+function useReviewedCellMarker(cellId: CellId | null) {
+  useEffect(() => {
+    if (cellId === null) {
+      return;
+    }
+    const cell = document.getElementById(HTMLCellId.create(cellId));
+    cell?.setAttribute("data-ai-review-selected", "");
+    return () => cell?.removeAttribute("data-ai-review-selected");
+  }, [cellId]);
+}

package/src/components/editor/chrome/wrapper/pending-ai-changes-badge.tsx

@@ -0,0 +1,25 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+import { useAtomValue } from "jotai";
+import { stagedAICellsAtom } from "@/core/ai/staged-cells";
+import { cn } from "@/utils/cn";
+
+export function PendingAIChangesBadge({ className }: { className?: string }) {
+  const count = useAtomValue(stagedAICellsAtom).size;
+  if (count === 0) {
+    return null;
+  }
+  return (
+    <span
+      className={cn(
+        "inline-flex h-4 min-w-4 items-center justify-center rounded-full border border-border bg-muted px-1 text-[10px] leading-none tabular-nums text-foreground",
+        className,
+      )}
+    >
+      <span aria-hidden={true}>{count}</span>
+      <span className="sr-only">
+        {count} pending AI {count === 1 ? "change" : "changes"}
+      </span>
+    </span>
+  );
+}

package/src/components/editor/chrome/wrapper/sidebar.tsx

@@ -7,10 +7,8 @@ import type { PropsWithChildren } from "react";
 import { useEffect, useMemo } from "react";
 import { ReorderableList } from "@/components/ui/reorderable-list";
 import { Tooltip } from "@/components/ui/tooltip";
-import {
-  cellErrorCount,
-  notebookQueuedOrRunningCountAtom,
-} from "@/core/cells/cells";
+import { cellErrorCount } from "@/core/cells/cells";
+import { notebookQueuedOrRunningCountAtom } from "@/core/executions";
 import { capabilitiesAtom } from "@/core/config/capabilities";
 import { aiEnabledAtom } from "@/core/config/config";
 import { cn } from "@/utils/cn";
@@ -22,6 +20,7 @@ import {
   PANELS,
   type PanelDescriptor,
 } from "../types";
+import { PendingAIChangesBadge } from "./pending-ai-changes-badge";
 import { PANEL_PRELOADERS } from "./lazy-panels";
 
 export const Sidebar: React.FC = () => {
@@ -150,6 +149,11 @@ export const Sidebar: React.FC = () => {
           >
             {panel.type === "errors" ? (
               <ErrorPanelIcon Icon={panel.Icon} />
+            ) : panel.type === "ai" ? (
+              <span className="relative">
+                {renderIcon(panel)}
+                <PendingAIChangesBadge className="absolute -top-2 -right-2" />
+              </span>
             ) : (
               renderIcon(panel)
             )}