@marimo-team/islands 0.25.3-dev30 → 0.25.3-dev301

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

@@ -1,24 +1,17 @@
 /* Copyright 2026 Marimo. All rights reserved. */
 
-import { useAtomValue, useStore } from "jotai";
-import { SparklesIcon } from "lucide-react";
-import { Tooltip } from "@/components/ui/tooltip";
+import { atom, useAtomValue, useStore } from "jotai";
+import { useMemo } from "react";
 import {
   type Edit,
-  stagedGenerationInProgressAtom,
+  stagedAICellsAtom,
   useStagedAICell,
   useStagedCells,
 } from "@/core/ai/staged-cells";
-import { getCellEditorView } from "@/core/cells/cells";
+import { chatGenerationInProgressAtom } from "@/core/ai/state";
 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 {
-  AI_EDIT_ACTION_LABELS,
-  CompletionActionsCellFooter,
-} from "../ai/completion-handlers";
-import { useRunCell } from "./useRunCells";
+import { AI_EDIT_ACTION_LABELS } from "../ai/completion-handlers";
+import type { ActionButton } from "../actions/types";
 
 const actionLabels: Record<Edit["type"], { accept: string; decline: string }> =
   {
@@ -27,134 +20,56 @@ const actionLabels: Record<Edit["type"], { accept: string; decline: string }> =
     delete_cell: { accept: "Delete cell", decline: "Keep cell" },
   };
 
-export const StagedAICellBackground: React.FC<{
-  cellId: CellId;
-  className?: string;
-}> = ({ cellId, className }) => {
-  const stagedCell = useStagedAICell(cellId);
-
-  if (!stagedCell) {
-    return null;
-  }
-
-  const cellClass =
-    stagedCell.type === "delete_cell"
-      ? "mo-ai-deleted-cell"
-      : "mo-ai-generated-cell";
-
-  return <div className={cn(cellClass, className)} />;
-};
-
-export const StagedAICellFooter: React.FC<{ cellId: CellId }> = ({
+export const StagedAIDeletionMarker: React.FC<{ cellId: CellId }> = ({
   cellId,
-}) => {
+}) =>
+  useStagedAICell(cellId)?.type === "delete_cell" ? (
+    <div className="mo-ai-deleted-cell" />
+  ) : null;
+
+export function useStagedAICellActions(
+  cellId: CellId | undefined,
+): ActionButton[] {
   const store = useStore();
-  const generationInProgress = useAtomValue(stagedGenerationInProgressAtom);
-  const stagedAiCell = useStagedAICell(cellId);
-  const runCell = useRunCell(cellId);
-
-  const { deleteStagedCell, removeStagedCell } = useStagedCells(store);
-
-  if (!stagedAiCell) {
-    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="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={generationInProgress}
-          onAccept={() => handleCompletion("accept")}
-          onDecline={() => handleCompletion("reject")}
-          size="xs"
-          runCell={isDeletion ? undefined : runCell}
-          acceptLabel={labels.accept}
-          declineLabel={labels.decline}
-        />
-      </div>
-    </div>
+  const edit = useStagedAICell(cellId);
+  // Every cell calls this; only cells with a pending edit care about the lock.
+  const isLocked = useAtomValue(
+    useMemo(
+      () =>
+        atom(
+          (get) =>
+            cellId !== undefined &&
+            get(stagedAICellsAtom).has(cellId) &&
+            get(chatGenerationInProgressAtom),
+        ),
+      [cellId],
+    ),
   );
-};
+  const { acceptStagedCell, rejectStagedCell } = useStagedCells();
 
-/**
- * Accept a staged cell and apply the changes.
- */
-export function acceptStagedCell(
-  cellId: CellId,
-  edit: Edit,
-  removeStagedCell: (cellId: CellId) => void,
-  deleteStagedCell: (cellId: CellId) => void,
-): void {
-  switch (edit.type) {
-    case "delete_cell":
-      // For delete cells, the cell is deleted when the completion is accepted
-      deleteStagedCell(cellId);
-      break;
-    default:
-      removeStagedCell(cellId);
-      break;
+  if (!edit || cellId === undefined) {
+    return [];
   }
-}
 
-/**
- * Reject a staged cell and revert the changes.
- */
-export function rejectStagedCell(
-  cellId: CellId,
-  edit: Edit,
-  removeStagedCell: (cellId: CellId) => void,
-  deleteStagedCell: (cellId: CellId) => void,
-): void {
-  switch (edit.type) {
-    case "update_cell": {
-      // Revert cell code
-      const editorView = getCellEditorView(cellId);
-      if (!editorView) {
-        Logger.error("Editor for this cell not found", { cellId });
-        break;
-      }
-
-      updateEditorCodeFromPython(editorView, edit.previousCode);
-      removeStagedCell(cellId);
-      break;
-    }
-    case "add_cell":
-      // Delete the cell since it's newly created
-      deleteStagedCell(cellId);
-      break;
-    case "delete_cell":
-      // Just remove the deletion marker - cell stays in notebook
-      removeStagedCell(cellId);
-      break;
-  }
+  const labels = actionLabels[edit.type];
+  return [
+    {
+      label: labels.accept,
+      disabled: isLocked,
+      handle: () => {
+        if (!store.get(chatGenerationInProgressAtom)) {
+          acceptStagedCell(cellId, edit);
+        }
+      },
+    },
+    {
+      label: labels.decline,
+      disabled: isLocked,
+      handle: () => {
+        if (!store.get(chatGenerationInProgressAtom)) {
+          rejectStagedCell(cellId, edit);
+        }
+      },
+    },
+  ];
 }

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

@@ -3,8 +3,8 @@ import { SquareIcon } from "lucide-react";
 import type { JSX } from "react";
 import { useRequestClient } from "@/core/network/requests";
 import type { RuntimeState } from "@/core/network/types";
-import { isAppInteractionDisabled } from "@/core/websocket/connection-utils";
-import type { ConnectionState } from "@/core/websocket/types";
+import { isAppInteractionDisabled } from "@/core/session/connection-utils";
+import type { ConnectionState } from "@/core/session/types";
 import { Functions } from "@/utils/functions";
 import { renderShortcut } from "../../shortcuts/renderShortcut";
 import { ToolbarItem } from "./toolbar";

package/src/components/editor/cell/__tests__/useRunCells.test.tsx

@@ -0,0 +1,169 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+import { act, renderHook } from "@testing-library/react";
+import { Provider } from "jotai";
+import type { ReactNode } from "react";
+import { afterEach, describe, expect, it, vi } from "vitest";
+import { MockExecutions } from "@/__mocks__/executions";
+import { MockNotebook } from "@/__mocks__/notebook";
+import { MockRequestClient } from "@/__mocks__/requests";
+import { cellId } from "@/__tests__/branded";
+import { notebookAtom } from "@/core/cells/cells";
+import {
+  cellExecutionInputsAtom,
+  cellExecutionView,
+  notebookQueuedOrRunningCountAtom,
+} from "@/core/executions";
+import { requestClientAtom } from "@/core/network/requests";
+import type { DocumentModel } from "@/core/notebook/document";
+import {
+  cell,
+  document as documentOf,
+  withCells,
+} from "@/core/notebook/document/__tests__/fixtures";
+import { activeDocumentModelAtom } from "@/core/session/active-notebook";
+import { store } from "@/core/state/jotai";
+import { runCells, useRunStaleCells } from "../useRunCells";
+
+const ENABLED = cellId("enabled");
+const DISABLED = cellId("disabled");
+
+describe("runCells", () => {
+  it("sends a disabled cell to the kernel without reading it as queued", async () => {
+    MockExecutions.seed(store, []);
+    const notebook = MockNotebook.notebookState({
+      cellData: {
+        [ENABLED]: { code: "x = 1" },
+        [DISABLED]: { code: "y = 2", config: { disabled: true } },
+      },
+    });
+    const sendRun = vi.fn(async () => null);
+
+    await runCells({
+      cellIds: [ENABLED, DISABLED],
+      sendRun,
+      prepareForRun: vi.fn(),
+      notebook,
+    });
+
+    expect(sendRun).toHaveBeenCalledWith({
+      cellIds: [ENABLED, DISABLED],
+      codes: ["x = 1", "y = 2"],
+    });
+    const inputs = store.get(cellExecutionInputsAtom);
+    expect(cellExecutionView(inputs, ENABLED).status).toBe("queued");
+    expect(cellExecutionView(inputs, DISABLED).status).toBe("idle");
+    expect(store.get(notebookQueuedOrRunningCountAtom)).toBe(1);
+  });
+});
+
+describe("useRunStaleCells", () => {
+  const C1 = cellId("c1");
+  const C2 = cellId("c2");
+
+  afterEach(() => {
+    store.set(activeDocumentModelAtom, null);
+  });
+
+  it("flushes pending edits before choosing the stale cells", async () => {
+    const document = documentOf();
+    const replica = MockExecutions.seed(
+      store,
+      [
+        MockExecutions.completed(C1, { code_fingerprint: "c1-hash" }),
+        MockExecutions.completed(C2, { code_fingerprint: "c2-hash" }),
+      ],
+      null,
+      document,
+    );
+    store.set(
+      notebookAtom,
+      MockNotebook.notebookState({
+        cellData: { [C1]: { code: "x = 2" }, [C2]: { code: "y = 2" } },
+      }),
+    );
+    const client = MockRequestClient.create();
+    store.set(requestClientAtom, client);
+    // The edit to c1 is still in flight: the host confirms it on flush.
+    const flush = vi.fn(async () => {
+      await new Promise((resolve) => setTimeout(resolve, 0));
+      store.set(
+        replica.documentAtom,
+        withCells(document, {
+          [C1]: cell(C1, "x = 2", { code_fingerprint: "c1-hash-2" }),
+        }),
+      );
+    });
+    store.set(activeDocumentModelAtom, { flush } as unknown as DocumentModel);
+    const wrapper = ({ children }: { children: ReactNode }) => (
+      <Provider store={store}>{children}</Provider>
+    );
+    const { result } = renderHook(useRunStaleCells, { wrapper });
+
+    await act(() => result.current());
+
+    expect(flush).toHaveBeenCalledOnce();
+    expect(client.sendRun).toHaveBeenCalledWith({
+      cellIds: [C1],
+      codes: ["x = 2"],
+    });
+  });
+
+  it("does nothing when the notebook changed while the flush was pending", async () => {
+    MockExecutions.seed(
+      store,
+      [MockExecutions.completed(C1, { code_fingerprint: "stale" })],
+      null,
+      documentOf(),
+    );
+    store.set(
+      notebookAtom,
+      MockNotebook.notebookState({ cellData: { [C1]: { code: "x = 1" } } }),
+    );
+    const client = MockRequestClient.create();
+    store.set(requestClientAtom, client);
+    const flush = vi.fn(async () => {
+      store.set(activeDocumentModelAtom, null);
+      throw new Error("aborted");
+    });
+    store.set(activeDocumentModelAtom, { flush } as unknown as DocumentModel);
+    const wrapper = ({ children }: { children: ReactNode }) => (
+      <Provider store={store}>{children}</Provider>
+    );
+    const { result } = renderHook(useRunStaleCells, { wrapper });
+
+    await act(() => result.current());
+
+    expect(client.sendRun).not.toHaveBeenCalled();
+  });
+
+  it("still runs what is stale when the flush fails", async () => {
+    MockExecutions.seed(
+      store,
+      [MockExecutions.completed(C1, { code_fingerprint: "stale" })],
+      null,
+      documentOf(),
+    );
+    store.set(
+      notebookAtom,
+      MockNotebook.notebookState({ cellData: { [C1]: { code: "x = 1" } } }),
+    );
+    const client = MockRequestClient.create();
+    store.set(requestClientAtom, client);
+    const flush = vi.fn(async () => {
+      throw new Error("offline");
+    });
+    store.set(activeDocumentModelAtom, { flush } as unknown as DocumentModel);
+    const wrapper = ({ children }: { children: ReactNode }) => (
+      <Provider store={store}>{children}</Provider>
+    );
+    const { result } = renderHook(useRunStaleCells, { wrapper });
+
+    await act(() => result.current());
+
+    expect(client.sendRun).toHaveBeenCalledWith({
+      cellIds: [C1],
+      codes: ["x = 1"],
+    });
+  });
+});

package/src/components/editor/cell/__tests__/useShouldShowInterrupt.test.ts

@@ -0,0 +1,39 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+import { act, renderHook } from "@testing-library/react";
+import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
+import { useShouldShowInterrupt } from "../useShouldShowInterrupt";
+
+describe("useShouldShowInterrupt", () => {
+  beforeEach(() => vi.useFakeTimers());
+  afterEach(() => vi.useRealTimers());
+
+  it("waits for each run and hides immediately when it stops", () => {
+    const { result, rerender } = renderHook(useShouldShowInterrupt, {
+      initialProps: false,
+    });
+    expect(result.current).toBe(false);
+    rerender(true);
+    expect(result.current).toBe(false);
+    act(() => vi.advanceTimersByTime(200));
+    expect(result.current).toBe(true);
+    rerender(false);
+    expect(result.current).toBe(false);
+    rerender(true);
+    expect(result.current).toBe(false);
+    act(() => vi.advanceTimersByTime(199));
+    expect(result.current).toBe(false);
+    act(() => vi.advanceTimersByTime(1));
+    expect(result.current).toBe(true);
+  });
+
+  it("cancels the delay for a short run", () => {
+    const { result, rerender } = renderHook(useShouldShowInterrupt, {
+      initialProps: true,
+    });
+    act(() => vi.advanceTimersByTime(100));
+    rerender(false);
+    act(() => vi.advanceTimersByTime(200));
+    expect(result.current).toBe(false);
+  });
+});

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

@@ -60,6 +60,7 @@ const CellActionsDropdownInternal = (
   ref: React.Ref<CellActionsDropdownHandle>,
 ) => {
   const [open, setOpen] = useState(false);
+  const pendingAction = useRef<(() => void) | null>(null);
   const closePopover = () => setOpen(false);
   const buttonRef = useRef<HTMLButtonElement>(null);
   const actions = useCellActionButtons({ cell: props, closePopover });
@@ -107,6 +108,16 @@ const CellActionsDropdownInternal = (
       className="w-[300px] p-0 pt-1 overflow-auto"
       scrollable={true}
       {...restoreFocus}
+      onCloseAutoFocus={(event) => {
+        restoreFocus.onCloseAutoFocus(event);
+        const action = pendingAction.current;
+        pendingAction.current = null;
+        if (action) {
+          // Restore focus before the action moves it or hides the editor.
+          event.preventDefault();
+          action();
+        }
+      }}
     >
       <Command>
         <CommandInput
@@ -153,7 +164,7 @@ const CellActionsDropdownInternal = (
                         if (action.disableClick || action.disabled) {
                           return;
                         }
-                        action.handle();
+                        pendingAction.current = action.handle;
                         setOpen(false);
                       }}
                       variant={action.variant}

package/src/components/editor/cell/cell-context-menu.tsx

@@ -23,6 +23,7 @@ import { toast } from "@/components/ui/use-toast";
 import { useCellData, useCellRuntime } from "@/core/cells/cells";
 import { CellOutputId } from "@/core/cells/ids";
 import { isOutputEmpty } from "@/core/cells/outputs";
+import { useCellExecution } from "@/core/executions";
 import { goToDefinitionWithLspFallback } from "@/core/codemirror/go-to-definition/utils";
 import { sendToPanelManager } from "@/core/vscode/vscode-bindings";
 import { copyImageToClipboard, copyToClipboard } from "@/utils/copy";
@@ -54,12 +55,14 @@ export const CellActionsContextMenu = ({
 }: Props) => {
   const cellData = useCellData(cellId);
   const cellRuntime = useCellRuntime(cellId);
+  const execution = useCellExecution(cellId);
   const actions = useCellActionButtons({
     cell: {
       cellId: cellId,
       name: cellData.name,
       config: cellData.config,
-      status: cellRuntime.status,
+      status: execution.status,
+      blocked: execution.blocked,
       hasOutput: !isOutputEmpty(cellRuntime.output),
       hasConsoleOutput: cellRuntime.consoleOutputs.length > 0,
       getEditorView,

package/src/components/editor/cell/code/__tests__/language-toggle.test.tsx

@@ -0,0 +1,54 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+import { EditorState } from "@codemirror/state";
+import { EditorView } from "@codemirror/view";
+import { fireEvent, render, screen } from "@testing-library/react";
+import { createRef } from "react";
+import { afterEach, describe, expect, it, vi } from "vitest";
+import { TooltipProvider } from "@/components/ui/tooltip";
+import { switchLanguage } from "@/core/codemirror/language/extension";
+import { LanguageToggle } from "../language-toggle";
+
+vi.mock("@/core/codemirror/language/extension", () => ({
+  switchLanguage: vi.fn(),
+}));
+
+describe("LanguageToggle", () => {
+  afterEach(() => {
+    vi.mocked(switchLanguage).mockClear();
+  });
+
+  it("uses the editor mounted after the button rendered", () => {
+    const editorViewRef = createRef<EditorView>();
+    const onAfterToggle = vi.fn();
+    render(
+      <TooltipProvider>
+        <LanguageToggle
+          editorViewRef={editorViewRef}
+          currentLanguageAdapter="python"
+          canSwitchToLanguage={true}
+          toType="markdown"
+          displayName="Markdown"
+          icon="Markdown"
+          onAfterToggle={onAfterToggle}
+        />
+      </TooltipProvider>,
+    );
+
+    fireEvent.click(screen.getByRole("button", { name: "Markdown" }));
+    expect(switchLanguage).not.toHaveBeenCalled();
+    expect(onAfterToggle).not.toHaveBeenCalled();
+
+    const view = new EditorView({ state: EditorState.create() });
+    editorViewRef.current = view;
+    try {
+      fireEvent.click(screen.getByRole("button", { name: "Markdown" }));
+      expect(switchLanguage).toHaveBeenCalledWith(view, {
+        language: "markdown",
+      });
+      expect(onAfterToggle).toHaveBeenCalledOnce();
+    } finally {
+      view.destroy();
+    }
+  });
+});