@marimo-team/frontend 0.24.1-dev53 → 0.24.1-dev54

package/src/components/editor/ai/completion-handlers.tsx

@@ -7,6 +7,11 @@ import { Button } from "@/components/ui/button";
 import { Tooltip } from "@/components/ui/tooltip";
 import { isPlatformMac } from "@/core/hotkeys/shortcuts";
 
+export const AI_EDIT_ACTION_LABELS = {
+  accept: "Keep change",
+  decline: "Revert change",
+} as const;
+
 /**
  * Common keyboard shortcut handlers for AI completions
  */
@@ -48,16 +53,31 @@ export const CompletionActionsCellFooter: React.FC<{
   size?: "xs" | "sm";
   multipleCompletions?: boolean;
   runCell?: () => void;
-}> = ({ isLoading, onAccept, onDecline, runCell }) => {
+  acceptLabel?: string;
+  declineLabel?: string;
+}> = ({
+  isLoading,
+  onAccept,
+  onDecline,
+  runCell,
+  acceptLabel,
+  declineLabel,
+  size,
+}) => {
   return (
     <>
       <AcceptCompletionButton
         isLoading={isLoading}
         onAccept={onAccept}
-        size="xs"
+        size={size}
         runCell={runCell}
+        label={acceptLabel}
+      />
+      <RejectCompletionButton
+        onDecline={onDecline}
+        size={size}
+        label={declineLabel}
       />
-      <RejectCompletionButton onDecline={onDecline} size="xs" />
     </>
   );
 };
@@ -72,6 +92,7 @@ export const AcceptCompletionButton: React.FC<{
   borderless?: boolean;
   acceptShortcut?: string;
   runCell?: () => void;
+  label?: string;
 }> = ({
   isLoading,
   onAccept,
@@ -82,6 +103,7 @@ export const AcceptCompletionButton: React.FC<{
   runCell,
   playButtonStyles,
   borderless = false,
+  label,
 }) => {
   const handleAcceptAndRun = () => {
     onAccept();
@@ -90,7 +112,7 @@ export const AcceptCompletionButton: React.FC<{
     }
   };
 
-  const text = multipleCompletions ? "Accept all" : "Accept";
+  const text = label ?? (multipleCompletions ? "Accept all" : "Accept");
 
   const baseClasses = `h-6 text-(--grass-11) bg-(--grass-3)/60
     hover:bg-(--grass-3) dark:bg-(--grass-4)/80 dark:hover:bg-(--grass-3) font-semibold
@@ -113,11 +135,7 @@ export const AcceptCompletionButton: React.FC<{
           )}
         </Button>
         <Tooltip
-          content={
-            multipleCompletions
-              ? "Accept and run all cells"
-              : "Accept and run cell"
-          }
+          content={`${text} and run ${multipleCompletions ? "all cells" : "cell"}`}
         >
           <Button
             variant="text"
@@ -156,6 +174,7 @@ export const RejectCompletionButton: React.FC<{
   className?: string;
   borderless?: boolean;
   declineShortcut?: string;
+  label?: string;
 }> = ({
   onDecline,
   multipleCompletions = false,
@@ -163,6 +182,7 @@ export const RejectCompletionButton: React.FC<{
   className,
   declineShortcut,
   borderless = false,
+  label,
 }) => {
   return (
     <Button
@@ -174,7 +194,7 @@ export const RejectCompletionButton: React.FC<{
     active:bg-(--red-5) dark:active:bg-(--red-4)
     border-(--red-6) border hover:shadow-xs ${borderless && "border-none rounded-md"} ${className}`}
     >
-      Reject{multipleCompletions && " all"}
+      {label ?? `Reject${multipleCompletions ? " all" : ""}`}
       {declineShortcut && (
         <MinimalHotkeys className="ml-1 inline" shortcut={declineShortcut} />
       )}

package/src/components/editor/cell/StagedAICell.tsx

@@ -1,9 +1,11 @@
 /* Copyright 2026 Marimo. All rights reserved. */
 
-import { useAtomValue, useStore } from "jotai";
+import { useStore } from "jotai";
+import { SparklesIcon } from "lucide-react";
+import { Tooltip } from "@/components/ui/tooltip";
 import {
   type Edit,
-  stagedAICellsAtom,
+  useStagedAICell,
   useStagedCells,
 } from "@/core/ai/staged-cells";
 import { getCellEditorView } from "@/core/cells/cells";
@@ -11,15 +13,24 @@ import type { CellId } from "@/core/cells/ids";
 import { updateEditorCodeFromPython } from "@/core/codemirror/language/utils";
 import { cn } from "@/utils/cn";
 import { Logger } from "@/utils/Logger";
-import { CompletionActionsCellFooter } from "../ai/completion-handlers";
+import {
+  AI_EDIT_ACTION_LABELS,
+  CompletionActionsCellFooter,
+} from "../ai/completion-handlers";
 import { useRunCell } from "./useRunCells";
 
+const actionLabels: Record<Edit["type"], { accept: string; decline: string }> =
+  {
+    add_cell: { accept: "Keep cell", decline: "Discard cell" },
+    update_cell: AI_EDIT_ACTION_LABELS,
+    delete_cell: { accept: "Delete cell", decline: "Keep cell" },
+  };
+
 export const StagedAICellBackground: React.FC<{
   cellId: CellId;
   className?: string;
 }> = ({ cellId, className }) => {
-  const stagedAICells = useAtomValue(stagedAICellsAtom);
-  const stagedCell = stagedAICells.get(cellId);
+  const stagedCell = useStagedAICell(cellId);
 
   if (!stagedCell) {
     return null;
@@ -37,8 +48,7 @@ export const StagedAICellFooter: React.FC<{ cellId: CellId }> = ({
   cellId,
 }) => {
   const store = useStore();
-  const stagedAICells = useAtomValue(stagedAICellsAtom);
-  const stagedAiCell = stagedAICells.get(cellId);
+  const stagedAiCell = useStagedAICell(cellId);
   const runCell = useRunCell(cellId);
 
   const { deleteStagedCell, removeStagedCell } = useStagedCells(store);
@@ -47,21 +57,49 @@ export const StagedAICellFooter: React.FC<{ cellId: CellId }> = ({
     return null;
   }
 
+  const isDeletion = stagedAiCell.type === "delete_cell";
+  const labels = actionLabels[stagedAiCell.type];
+
   const handleCompletion = (type: "accept" | "reject") => {
     const completionFunc =
       type === "accept" ? acceptStagedCell : rejectStagedCell;
     completionFunc(cellId, stagedAiCell, removeStagedCell, deleteStagedCell);
   };
 
+  const tooltipContent =
+    stagedAiCell.type === "add_cell"
+      ? "AI-generated cell"
+      : isDeletion
+        ? "AI suggests deleting this cell"
+        : "AI changed this cell";
+
   return (
-    <div className="flex items-center justify-end gap-1.5 w-full pb-1 pt-2">
-      <CompletionActionsCellFooter
-        isLoading={false}
-        onAccept={() => handleCompletion("accept")}
-        onDecline={() => handleCompletion("reject")}
-        size="xs"
-        runCell={runCell}
-      />
+    <div className="mo-ai-cell-footer">
+      <div className="flex items-center gap-1.5 text-xs text-muted-foreground">
+        <Tooltip content={tooltipContent}>
+          <span
+            className="inline-flex items-center gap-1.5"
+            aria-label={tooltipContent}
+          >
+            <SparklesIcon
+              aria-hidden="true"
+              className="size-3.5 text-(--blue-11)"
+            />
+            <span>AI</span>
+          </span>
+        </Tooltip>
+      </div>
+      <div className="flex items-center gap-1.5">
+        <CompletionActionsCellFooter
+          isLoading={false}
+          onAccept={() => handleCompletion("accept")}
+          onDecline={() => handleCompletion("reject")}
+          size="xs"
+          runCell={isDeletion ? undefined : runCell}
+          acceptLabel={labels.accept}
+          declineLabel={labels.decline}
+        />
+      </div>
     </div>
   );
 };

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

@@ -56,7 +56,10 @@ export const PendingAICells: React.FC = () => {
   };
 
   const runAllCells = () => {
-    runCell(listStagedCells);
+    const cellsToRun = listStagedCells.filter(
+      (cellId) => stagedAiCells.get(cellId)?.type !== "delete_cell",
+    );
+    runCell(cellsToRun);
   };
 
   const cyanShadow = "shadow-[0_0_6px_0_#00A2C733]";
@@ -91,6 +94,7 @@ export const PendingAICells: React.FC = () => {
           multipleCompletions={true}
           onAccept={acceptAllCompletions}
           isLoading={false}
+          label="Keep all"
           size="xs"
           buttonStyles="h-6.5"
           playButtonStyles="h-6.5"
@@ -100,6 +104,7 @@ export const PendingAICells: React.FC = () => {
           multipleCompletions={true}
           onDecline={rejectAllCompletions}
           size="xs"
+          label="Discard all"
           className="h-6.5"
         />
       </div>

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

@@ -839,8 +839,8 @@ const EditableCellComponent = ({
               />
             )}
             <PendingDeleteConfirmation cellId={cellId} />
+            {!isPresenting && <StagedAICellFooter cellId={cellId} />}
           </div>
-          {!isPresenting && <StagedAICellFooter cellId={cellId} />}
           {isCollapsed && !isPresenting && (
             <CollapsedCellBanner
               onClick={() => actions.expandCell({ cellId })}
@@ -1327,8 +1327,8 @@ const SetupCellComponent = ({
               cellId={cellId}
               debuggerActive={cellRuntime.debuggerActive}
             />
+            <StagedAICellFooter cellId={cellId} />
           </div>
-          <StagedAICellFooter cellId={cellId} />
         </div>
       </CellActionsContextMenu>
     </TooltipProvider>

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

@@ -1,7 +1,8 @@
 /* Copyright 2026 Marimo. All rights reserved. */
 
-import { renderHook } from "@testing-library/react";
-import { getDefaultStore } from "jotai";
+import { act, renderHook } from "@testing-library/react";
+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";
@@ -9,6 +10,7 @@ import { updateEditorCodeFromPython } from "../../codemirror/language/utils";
 import {
   type StagedAICells,
   stagedAICellsAtom,
+  useStagedAICell,
   useStagedCells,
   visibleForTesting,
 } from "../staged-cells";
@@ -149,6 +151,16 @@ describe("staged-cells", () => {
       expect(newState.has(cellId2)).toBe(true);
     });
 
+    it("should preserve state when removing a cell that is not staged", () => {
+      const state = new Map([[cellId1, { type: "add_cell" as const }]]);
+      const newState = reducer(state, {
+        type: "removeStagedCell",
+        payload: cellId2,
+      });
+
+      expect(newState).toBe(state);
+    });
+
     it("should clear all cells", () => {
       const state = new Map([
         [cellId1, { type: "add_cell" as const }],
@@ -206,6 +218,37 @@ describe("staged-cells", () => {
       const state = store.get(stagedAICellsAtom);
       expect(state).toEqual(new Map());
     });
+
+    it("should not rerender a cell when another cell's staging changes", () => {
+      let renderCount = 0;
+      const wrapper = ({ children }: { children: ReactNode }) =>
+        createElement(Provider, { store }, children);
+
+      renderHook(
+        () => {
+          renderCount++;
+          return useStagedAICell(cellId2);
+        },
+        { wrapper },
+      );
+      const initialRenderCount = renderCount;
+
+      act(() => {
+        store.set(
+          stagedAICellsAtom,
+          new Map([[cellId1, { type: "add_cell" }]]),
+        );
+      });
+      expect(renderCount).toBe(initialRenderCount);
+
+      act(() => {
+        store.set(
+          stagedAICellsAtom,
+          new Map([[cellId2, { type: "add_cell" }]]),
+        );
+      });
+      expect(renderCount).toBeGreaterThan(initialRenderCount);
+    });
   });
 
   describe("useStagedCells hook", () => {

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

@@ -1,7 +1,9 @@
 /* Copyright 2026 Marimo. All rights reserved. */
 
 import type { UIMessageChunk } from "ai";
-import { useRef } from "react";
+import { useAtomValue } from "jotai";
+import { selectAtom } from "jotai/utils";
+import { useMemo, useRef } from "react";
 import {
   type AiCompletion,
   codeToCells,
@@ -51,6 +53,9 @@ const {
     return new Map([...state, [cellId, edit]]);
   },
   removeStagedCell: (state, cellId: CellId) => {
+    if (!state.has(cellId)) {
+      return state;
+    }
     const newState = new Map(state);
     newState.delete(cellId);
     return newState;
@@ -66,6 +71,15 @@ export {
   reducer as stagedAICellsReducer,
 };
 
+export function useStagedAICell(cellId: CellId): Edit | undefined {
+  const stagedCellAtom = useMemo(
+    () =>
+      selectAtom(stagedAICellsAtom, (cells) => cells.get(cellId), Object.is),
+    [cellId],
+  );
+  return useAtomValue(stagedCellAtom);
+}
+
 interface UpdateStagedCellAction {
   cellId: CellId;
   code: string;

package/src/css/app/Cell.css

@@ -27,25 +27,14 @@
     outline: none;
   }
 
-  &:has(.mo-ai-generated-cell)::before {
-    content: "";
-    position: absolute;
-    inset: 0;
-    border-radius: 10px;
-    pointer-events: none;
-    opacity: 0.5;
-    box-shadow: 0px 0px 10px 0px var(--cyan-9);
-    z-index: 1;
-  }
-
-  &:has(.mo-ai-generated-cell.mo-ai-setup-cell)::before {
-    border-radius: 2px;
+  &:has(.mo-ai-generated-cell) {
+    border-color: var(--blue-7);
   }
 
   /* Styling for cells marked for deletion by AI */
   &:has(.mo-ai-deleted-cell) {
-    opacity: 0.7;
     position: relative;
+    border-color: var(--red-7);
 
     .output-area,
     .cm-gutters,
@@ -53,18 +42,6 @@
       background-color: var(--red-2);
     }
 
-    /* Add a red glow to indicate deletion */
-    &::before {
-      content: "";
-      position: absolute;
-      inset: 0;
-      border-radius: 10px;
-      pointer-events: none;
-      opacity: 0.6;
-      box-shadow: 0px 0px 10px 0px var(--red-9);
-      z-index: 1;
-    }
-
     /* Add strikethrough effect to the code */
     .cm-content {
       text-decoration: line-through;
@@ -326,6 +303,20 @@
   @apply bg-background;
 }
 
+.mo-ai-cell-footer {
+  display: flex;
+  align-items: center;
+  justify-content: space-between;
+  gap: 0.75rem;
+  width: 100%;
+  padding: 0.375rem 0.5rem;
+  background: var(--gray-2);
+}
+
+.dark .mo-ai-cell-footer {
+  background: var(--gray-3);
+}
+
 .dark .marimo-cell {
   @apply border-border;
 
@@ -364,6 +355,14 @@
   }
 }
 
+.dark .marimo-cell:has(.mo-ai-generated-cell) {
+  border-color: var(--blue-8);
+}
+
+.dark .marimo-cell:has(.mo-ai-deleted-cell) {
+  border-color: var(--red-8);
+}
+
 #cell-setup,
 #cell-setup .cm-editor,
 #cell-setup .cm-gutter {