@marimo-team/islands 0.25.3-dev23 → 0.25.3-dev230

package/src/components/editor/actions/useConfigActions.tsx

@@ -8,7 +8,7 @@ import type { ActionButton } from "./types";
 
 export function useConfigActions() {
   const [config, setConfig] = useResolvedMarimoConfig();
-  const [appConfig, setAppConfig] = useAppConfig();
+  const appConfig = useAppConfig();
   const { saveAppConfig, saveUserConfig } = useRequestClient();
 
   const handleUserConfig = async (values: Partial<UserConfig>) => {
@@ -17,10 +17,9 @@ export function useConfigActions() {
     });
   };
 
+  // The host document carries the saved config back.
   const handleAppConfig = async (values: AppConfig) => {
-    await saveAppConfig({ config: values }).then(() => {
-      setAppConfig(values);
-    });
+    await saveAppConfig({ config: values });
   };
 
   const actions: ActionButton[] = [

package/src/components/editor/actions/useNotebookActions.tsx

@@ -62,22 +62,23 @@ import {
   getNotebook,
   hasDisabledCellsAtom,
   undoLabelAtom,
-  useCellActions,
 } from "@/core/cells/cells";
+import { useCellActions } from "@/core/cells/actions";
+import { useNotebookDocumentEdit } from "@/core/notebook/document";
+import { canEditAtom } from "@/core/notebook";
 import { disabledCellIds } from "@/core/cells/utils";
 import { capabilitiesAtom } from "@/core/config/capabilities";
 import { aiEnabledAtom, useResolvedMarimoConfig } from "@/core/config/config";
 import { Constants } from "@/core/constants";
 import { useLayoutActions, useLayoutState } from "@/core/layout/state";
 import { useTogglePresenting } from "@/core/layout/useTogglePresenting";
-import { kioskModeAtom, viewStateAtom } from "@/core/mode";
+import { viewStateAtom } from "@/core/mode";
 import { useRequestClient } from "@/core/network/requests";
 import { useFilename } from "@/core/saving/filename";
 import { createShareableLink } from "@/core/wasm/share";
 import { isWasm } from "@/core/wasm/utils";
 import { useDetectedDataSources } from "@/hooks/useDataSourceDiscovery";
 import { copyToClipboard } from "@/utils/copy";
-import { Objects } from "@/utils/objects";
 import { Strings } from "@/utils/strings";
 import { newNotebookURL } from "@/utils/urls";
 import { useRunAllCells } from "../cell/useRunCells";
@@ -119,7 +120,7 @@ export function useNotebookActions({
   const { toggleApplication } = useChromeActions();
   const { selectedPanel } = useChromeState();
   const [viewState] = useAtom(viewStateAtom);
-  const kioskMode = useAtomValue(kioskModeAtom);
+  const canEdit = useAtomValue(canEditAtom);
   const setCodeVisibility = useSetCodeVisibility();
   const setOutputsExpanded = useSetOutputsExpanded();
   const [resolvedConfig] = useResolvedMarimoConfig();
@@ -127,12 +128,11 @@ export function useNotebookActions({
   const aiEnabled = useAtomValue(aiEnabledAtom);
 
   const {
-    updateCellConfig,
     undoDeleteCell,
     clearAllCellOutputs,
-    addSetupCellIfDoesntExist,
     collapseAllCells,
     expandAllCells,
+    addSetupCell,
   } = useCellActions();
   const restartKernel = useRestartKernel();
   const runAllCells = useRunAllCells();
@@ -142,7 +142,8 @@ export function useNotebookActions({
   const { handleClick: openSettings } = useOpenSettingsToTab();
   const setKeyboardShortcutsOpen = useSetAtom(keyboardShortcutsAtom);
   const setExportOptions = useSetAtom(exportOptionsAtom);
-  const { readCode, saveCellConfig } = useRequestClient();
+  const { readCode } = useRequestClient();
+  const edit = useNotebookDocumentEdit();
   const detectedDatabaseSources = useDetectedDataSources("database");
   const detectedStorageSources = useDetectedDataSources("storage");
   const addDetectedDataSource = useAddDetectedDataSource();
@@ -301,7 +302,7 @@ export function useNotebookActions({
 
     {
       icon: <SparklesIcon size={14} strokeWidth={1.5} />,
-      label: "Connect your agent",
+      label: "Pair an agent",
       hidden: isWasm(),
       handle: async () => {
         openModal(<PairWithAgentModal onClose={closeModal} />);
@@ -448,19 +449,14 @@ export function useNotebookActions({
     {
       icon: <ZapIcon size={14} strokeWidth={1.5} />,
       label: "Enable all cells",
-      hidden: !hasDisabledCells || kioskMode,
-      handle: async () => {
-        const notebook = getNotebook();
-        const ids = disabledCellIds(notebook);
-        const newConfigs = Objects.fromEntries(
-          ids.map((cellId) => [cellId, { disabled: false }]),
-        );
-        // send to BE
-        await saveCellConfig({ configs: newConfigs });
-        // update on FE
-        for (const cellId of ids) {
-          updateCellConfig({ cellId, config: { disabled: false } });
-        }
+      hidden: !hasDisabledCells || !canEdit,
+      handle: () => {
+        const ids = disabledCellIds(getNotebook());
+        edit((tx) => {
+          for (const cellId of ids) {
+            tx.updateConfig(cellId, { disabled: false });
+          }
+        });
       },
     },
 
@@ -469,7 +465,7 @@ export function useNotebookActions({
       icon: <DiamondPlusIcon size={14} strokeWidth={1.5} />,
       label: "Add setup cell",
       handle: () => {
-        addSetupCellIfDoesntExist({});
+        addSetupCell();
       },
     },
     {
@@ -543,7 +539,7 @@ export function useNotebookActions({
     {
       icon: <Undo2Icon size={14} strokeWidth={1.5} />,
       label: undoLabel,
-      hidden: !canUndoDeletes || kioskMode,
+      hidden: !canUndoDeletes || !canEdit,
       handle: () => {
         undoDeleteCell();
       },

package/src/components/editor/actions/useRestartKernel.tsx

@@ -5,7 +5,7 @@ import { useImperativeModal } from "@/components/modal/ImperativeModal";
 import { AlertDialogDestructiveAction } from "@/components/ui/alert-dialog";
 import { connectionAtom } from "@/core/network/connection";
 import { useRequestClient } from "@/core/network/requests";
-import { ConnectionState } from "@/core/websocket/types";
+import { ConnectionState } from "@/core/session/types";
 import { reloadSafe } from "@/utils/reload-safe";
 
 export function useRestartKernel() {

package/src/components/editor/actions/useSetCodeVisibility.ts

@@ -1,10 +1,10 @@
 /* Copyright 2026 Marimo. All rights reserved. */
 
 import { useCallback } from "react";
-import { getNotebook, useCellActions } from "@/core/cells/cells";
+import { getNotebook } from "@/core/cells/cells";
 import type { CellId } from "@/core/cells/ids";
 import { MarkdownLanguageAdapter } from "@/core/codemirror/language/languages/markdown";
-import { useRequestClient } from "@/core/network/requests";
+import { useNotebookDocumentEdit } from "@/core/notebook/document";
 import type { CellConfig } from "@/core/network/types";
 import { Objects } from "@/utils/objects";
 
@@ -21,16 +21,15 @@ const cellCodeKind = (code: string): CellCodeKind =>
  * cell of the given kind.
  */
 export const useSetCodeVisibility = () => {
-  const { updateCellConfig } = useCellActions();
-  const { saveCellConfig } = useRequestClient();
+  const edit = useNotebookDocumentEdit();
 
   return useCallback(
-    async (hidden: boolean, kind: CellCodeKind) => {
+    (hidden: boolean, kind: CellCodeKind) => {
       const notebook = getNotebook();
 
       const newConfigs: Record<CellId, Partial<CellConfig>> = {};
 
-      for (const cellId of notebook.cellIds.inOrderIds) {
+      for (const cellId of notebook.cellIds) {
         const cell = notebook.cellData[cellId];
         if (cell === undefined || cell.config.hide_code === hidden) {
           continue;
@@ -46,12 +45,12 @@ export const useSetCodeVisibility = () => {
         return;
       }
 
-      await saveCellConfig({ configs: newConfigs });
-
-      for (const [cellId, config] of entries) {
-        updateCellConfig({ cellId, config });
-      }
+      edit((tx) => {
+        for (const [cellId, config] of entries) {
+          tx.updateConfig(cellId, config);
+        }
+      });
     },
-    [updateCellConfig, saveCellConfig],
+    [edit],
   );
 };

package/src/components/editor/actions/useSetOutputsExpanded.ts

@@ -1,14 +1,10 @@
 /* Copyright 2026 Marimo. All rights reserved. */
 
 import { useCallback } from "react";
-import {
-  getNotebook,
-  type NotebookState,
-  useCellActions,
-} from "@/core/cells/cells";
+import { getNotebook, type NotebookState } from "@/core/cells/cells";
 import type { CellId } from "@/core/cells/ids";
 import { isOutputEmpty } from "@/core/cells/outputs";
-import { useRequestClient } from "@/core/network/requests";
+import { useNotebookDocumentEdit } from "@/core/notebook/document";
 import type { CellConfig } from "@/core/network/types";
 import { Objects } from "@/utils/objects";
 
@@ -23,7 +19,7 @@ export function outputExpansionConfigs(
   expanded: boolean,
 ): Record<CellId, Partial<CellConfig>> {
   const configs: Record<CellId, Partial<CellConfig>> = {};
-  for (const cellId of notebook.cellIds.inOrderIds) {
+  for (const cellId of notebook.cellIds) {
     const cell = notebook.cellData[cellId];
     if (
       cell === undefined ||
@@ -44,24 +40,23 @@ export function outputExpansionConfigs(
  * every output to a fixed height, by setting each cell's `expand_output` config.
  */
 export const useSetOutputsExpanded = () => {
-  const { updateCellConfig } = useCellActions();
-  const { saveCellConfig } = useRequestClient();
+  const edit = useNotebookDocumentEdit();
 
   return useCallback(
-    async (expanded: boolean) => {
-      const newConfigs = outputExpansionConfigs(getNotebook(), expanded);
-
-      const entries = Objects.entries(newConfigs);
+    (expanded: boolean) => {
+      const entries = Objects.entries(
+        outputExpansionConfigs(getNotebook(), expanded),
+      );
       if (entries.length === 0) {
         return;
       }
 
-      await saveCellConfig({ configs: newConfigs });
-
-      for (const [cellId, config] of entries) {
-        updateCellConfig({ cellId, config });
-      }
+      edit((tx) => {
+        for (const [cellId, config] of entries) {
+          tx.updateConfig(cellId, config);
+        }
+      });
     },
-    [updateCellConfig, saveCellConfig],
+    [edit],
   );
 };

package/src/components/editor/ai/add-cell-with-ai.tsx

@@ -106,9 +106,9 @@ export const AddCellWithAI: React.FC<{
   const runtimeManager = useRuntimeManager();
 
   const inputRef = useRef<ReactCodeMirrorRef>(null);
-  const submissionController = useRef(
-    new StagedCellSubmissionController(),
-  ).current;
+  const [submissionController] = useState(
+    () => new StagedCellSubmissionController(),
+  );
 
   const fileInputRef = useRef<HTMLInputElement>(null);
   const { files, addFiles, removeFile } = useFileState();

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

@@ -408,6 +408,7 @@ export const AiCompletionEditor: React.FC<Props> = ({
               </div>
             </Tooltip>
             <Button
+              aria-label="Discard AI completion"
               data-testid="decline-completion-button"
               variant="text"
               size="icon"

package/src/components/editor/alerts/__tests__/connection-notice.test.tsx

@@ -11,7 +11,7 @@ import { AppConfigSchema } from "@/core/config/config-schema";
 import { kernelStartupErrorAtom } from "@/core/errors/state";
 import { connectionAtom } from "@/core/network/connection";
 import { sandboxAtom, sandboxSyncAtom } from "@/core/packages/sandbox-state";
-import { WebSocketClosedReason, ConnectionState } from "@/core/websocket/types";
+import { WebSocketClosedReason, ConnectionState } from "@/core/session/types";
 import { chromeAtom } from "../../chrome/state";
 import { runtimeDetailsExpandedAtom } from "../../chrome/panels/runtime-details-state";
 import { ConnectionNotice } from "../connection-notice";

package/src/components/editor/alerts/connection-notice.tsx

@@ -8,7 +8,7 @@ import type { AppConfig } from "@/core/config/config-schema";
 import { connectionAtom } from "@/core/network/connection";
 import { useConnectionNotice } from "@/core/network/useConnectionNotice";
 import { sandboxAtom } from "@/core/packages/sandbox-state";
-import { ConnectionState } from "@/core/websocket/types";
+import { ConnectionState } from "@/core/session/types";
 import { useInterval } from "@/hooks/useInterval";
 import { cn } from "@/utils/cn";
 import { runtimeDetailsExpandedAtom } from "../chrome/panels/runtime-details-state";

package/src/components/editor/alerts/startup-logs-alert.tsx

@@ -51,6 +51,7 @@ export const StartupLogsAlert: React.FC = () => {
             {isRunning ? "Initializing..." : "Initialized"}
           </span>
           <Button
+            aria-label="Dismiss startup logs"
             variant="text"
             data-testid="remove-startup-logs-button"
             size="icon"

package/src/components/editor/alerts/stdin-blocking-alert.tsx

@@ -12,7 +12,7 @@ const hasBlockingStdinAtom = atom((get) => {
   const notebook = get(notebookAtom);
 
   // Check each cell in order
-  for (const cellId of notebook.cellIds.inOrderIds) {
+  for (const cellId of notebook.cellIds) {
     // Check if the cell is idle and has unresolved stdin
     const runtime = notebook.cellRuntime[cellId];
     if (runtime.status === "idle") {

package/src/components/editor/app-container.tsx

@@ -8,7 +8,7 @@ import {
   type ConnectionStatus,
   WebSocketClosedReason,
   ConnectionState,
-} from "@/core/websocket/types";
+} from "@/core/session/types";
 import { cn } from "@/utils/cn";
 import { DynamicFavicon } from "./dynamic-favicon";
 import { StatusOverlay } from "./header/status";

package/src/components/editor/boundary/__tests__/ErrorBoundary.test.tsx

@@ -14,7 +14,7 @@ import { requestClientAtom } from "@/core/network/requests";
 import type { EnvironmentInfo } from "@/core/network/types";
 import { filenameAtom } from "@/core/saving/file-state";
 import { store } from "@/core/state/jotai";
-import { ConnectionState } from "@/core/websocket/types";
+import { ConnectionState } from "@/core/session/types";
 import { ErrorBoundary } from "../ErrorBoundary";
 
 vi.mock("@/utils/copy", () => ({

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

@@ -1,6 +1,6 @@
 /* Copyright 2026 Marimo. All rights reserved. */
 import { DatabaseIcon, DiamondPlusIcon, PlusIcon } from "lucide-react";
-import { useState } from "react";
+import { useRef, useState } from "react";
 import { Button } from "@/components/editor/inputs/Inputs";
 import { MinimalHotkeys } from "@/components/shortcuts/renderShortcut";
 import {
@@ -10,14 +10,14 @@ import {
   DropdownMenuTrigger,
 } from "@/components/ui/dropdown-menu";
 import { maybeAddMarimoImport } from "@/core/cells/add-missing-import";
-import { useCellActions } from "@/core/cells/cells";
+import { useCellActions } from "@/core/cells/actions";
 import { LanguageAdapters } from "@/core/codemirror/language/LanguageAdapters";
 import { MARKDOWN_INITIAL_HIDE_CODE } from "@/core/codemirror/language/languages/markdown";
 import {
   getConnectionTooltip,
   isAppInteractionDisabled,
-} from "@/core/websocket/connection-utils";
-import type { ConnectionState } from "@/core/websocket/types";
+} from "@/core/session/connection-utils";
+import type { ConnectionState } from "@/core/session/types";
 import { cn } from "@/utils/cn";
 import { Tooltip } from "../../ui/tooltip";
 import { MarkdownIcon, PythonIcon } from "./code/icons";
@@ -33,10 +33,11 @@ export const CreateCellButton = ({
   onClick: ((opts: { code: string; hideCode?: boolean }) => void) | undefined;
   oneClickShortcut: "shift" | "mod";
 }) => {
-  const { createNewCell, addSetupCellIfDoesntExist } = useCellActions();
+  const { createCell, addSetupCell } = useCellActions();
   const shortcut = `${oneClickShortcut}-Click`;
   const [open, setOpen] = useState(false);
-  const [justOpened, setJustOpened] = useState(false);
+  // Time the dropdown last opened, to ignore the click that opened it
+  const openedAt = useRef(0);
 
   const baseTooltipContent =
     getConnectionTooltip(connectionState) || tooltipContent;
@@ -61,7 +62,7 @@ export const CreateCellButton = ({
   // enabled; the user experience is confusing otherwise (how does the user
   // know they need to run import marimo as mo. first?).
   const addMarkdownCell = () => {
-    maybeAddMarimoImport({ autoInstantiate: true, createNewCell });
+    maybeAddMarimoImport({ autoInstantiate: true, createCell });
     onClick?.({
       code: LanguageAdapters.markdown.defaultCode,
       hideCode: MARKDOWN_INITIAL_HIDE_CODE,
@@ -69,23 +70,17 @@ export const CreateCellButton = ({
   };
 
   const addSQLCell = () => {
-    maybeAddMarimoImport({ autoInstantiate: true, createNewCell });
+    maybeAddMarimoImport({ autoInstantiate: true, createCell });
     onClick?.({ code: LanguageAdapters.sql.defaultCode });
   };
 
-  const addSetupCell = () => {
-    addSetupCellIfDoesntExist({});
-  };
-
   const renderIcon = (icon: React.ReactNode) => {
     return <div className="mr-3 text-muted-foreground">{icon}</div>;
   };
 
   const openDropdown = () => {
     setOpen(true);
-    setJustOpened(true);
-    // Allow interactions after a brief delay to prevent the dropdown items immediately being clicked
-    setTimeout(() => setJustOpened(false), 200);
+    openedAt.current = Date.now();
   };
 
   // We use onPointerDownCapture (not onPointerDown) to intercept events in
@@ -121,15 +116,14 @@ export const CreateCellButton = ({
   const handleOpenChange = (isOpen: boolean) => {
     setOpen(isOpen);
     if (isOpen) {
-      setJustOpened(true);
-      // Allow interactions after a brief delay
-      setTimeout(() => setJustOpened(false), 200);
+      openedAt.current = Date.now();
     }
   };
 
   const handleFirstItemClick = (e: React.MouseEvent) => {
-    // Hack to prevent the first item from being clicked when the dropdown is opened
-    if (justOpened) {
+    // Hack to prevent the first item from being clicked when the dropdown is
+    // opened: allow interactions only after a brief delay
+    if (Date.now() - openedAt.current < 200) {
       e.preventDefault();
       e.stopPropagation();
       return;
@@ -172,7 +166,7 @@ export const CreateCellButton = ({
           {renderIcon(<DatabaseIcon size={13} strokeWidth={1.5} />)}
           SQL cell
         </DropdownMenuItem>
-        <DropdownMenuItem onClick={addSetupCell}>
+        <DropdownMenuItem onClick={() => addSetupCell()}>
           {renderIcon(<DiamondPlusIcon size={13} strokeWidth={1.5} />)}
           Setup cell
         </DropdownMenuItem>

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

@@ -5,8 +5,8 @@ import type { RuntimeState } from "@/core/network/types";
 import {
   getConnectionTooltip,
   isAppInteractionDisabled,
-} from "@/core/websocket/connection-utils";
-import type { ConnectionState } from "@/core/websocket/types";
+} from "@/core/session/connection-utils";
+import type { ConnectionState } from "@/core/session/types";
 import { cn } from "@/utils/cn";
 import { Events } from "@/utils/events";
 import { Button } from "../../ui/button";
@@ -37,6 +37,7 @@ export const DeleteButton = (props: {
   return (
     <Tooltip content={tooltipMsg} usePortal={false}>
       <Button
+        aria-label="Delete cell"
         variant="ghost"
         size="icon"
         onClick={onClick}

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

@@ -5,8 +5,8 @@ import type { CellConfig, RuntimeState } from "@/core/network/types";
 import {
   getConnectionTooltip,
   isAppInteractionDisabled,
-} from "@/core/websocket/connection-utils";
-import type { ConnectionState } from "@/core/websocket/types";
+} from "@/core/session/connection-utils";
+import type { ConnectionState } from "@/core/session/types";
 import { renderShortcut } from "../../shortcuts/renderShortcut";
 import { ToolbarItem } from "./toolbar";
 

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

@@ -11,7 +11,7 @@ import {
 } from "@/core/ai/staged-cells";
 import { getCellEditorView } from "@/core/cells/cells";
 import type { CellId } from "@/core/cells/ids";
-import { updateEditorCodeFromPython } from "@/core/codemirror/language/utils";
+import { prepareCellCodeUpdate } from "@/core/codemirror/language/cell-code";
 import { cn } from "@/utils/cn";
 import { Logger } from "@/utils/Logger";
 import {
@@ -144,7 +144,15 @@ export function rejectStagedCell(
         break;
       }
 
-      updateEditorCodeFromPython(editorView, edit.previousCode);
+      const update = prepareCellCodeUpdate(
+        editorView.state,
+        edit.previousCode,
+        { initialize: true },
+      );
+
+      if (update) {
+        editorView.dispatch(update);
+      }
       removeStagedCell(cellId);
       break;
     }

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__/StagedAICell.test.tsx

@@ -14,9 +14,11 @@ vi.mock("@/components/editor/cell/useDeleteCell", () => ({
 
 vi.mock("@/core/cells/cells", () => ({
   getCellEditorView: vi.fn(),
+}));
+vi.mock("@/core/cells/actions", () => ({
   useCellActions: () => ({
-    createNewCell: vi.fn(),
-    updateCellCode: vi.fn(),
+    createCell: vi.fn(),
+    setCellCode: vi.fn(),
   }),
 }));
 

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}