@marimo-team/islands 0.25.3-dev21 → 0.25.3-dev211

package/src/components/editor/actions/__tests__/pair-with-agent.test.tsx

@@ -75,9 +75,7 @@ describe("PairWithAgent", () => {
       }),
     );
 
-    expect(
-      screen.getByRole("dialog", { name: "Pair with an agent" }),
-    ).toBeVisible();
+    expect(screen.getByRole("dialog", { name: "Pair an agent" })).toBeVisible();
   });
 
   it("opens the pairing dialog and preserves a passed click handler", () => {
@@ -88,9 +86,7 @@ describe("PairWithAgent", () => {
     fireEvent.click(screen.getByRole("button", { name: "Connect your agent" }));
 
     expect(onClick).toHaveBeenCalledOnce();
-    expect(
-      screen.getByRole("dialog", { name: "Pair with an agent" }),
-    ).toBeVisible();
+    expect(screen.getByRole("dialog", { name: "Pair an agent" })).toBeVisible();
   });
 
   it("hides pairing controls in WASM", () => {

package/src/components/editor/actions/__tests__/useCellActionButton.test.tsx

@@ -0,0 +1,113 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+import type { NotebookOutput } from "@marimo-team/client/protocol";
+import { renderHook } from "@testing-library/react";
+import { atom, createStore } from "jotai";
+import type { ReactNode } from "react";
+import { beforeEach, describe, expect, it, onTestFinished, vi } from "vitest";
+import { DocumentWrapper, mockDocumentModel } from "@/__mocks__/document";
+import { partial } from "@/__mocks__/partial";
+import { MockRequestClient } from "@/__mocks__/requests";
+import { ModalProvider } from "@/components/modal/ImperativeModal";
+import { toast } from "@/components/ui/use-toast";
+import { CellId } from "@/core/cells/ids";
+import { requestClientAtom } from "@/core/network/requests";
+import type { NotebookReplica } from "@/core/notebook/replica";
+import { activeNotebookReplicaAtom } from "@/core/session/active-notebook";
+import { copyToClipboard } from "@/utils/copy";
+import type { ActionButton } from "../types";
+import { useCellActionButtons } from "../useCellActionButton";
+
+vi.mock("@/utils/copy", () => ({ copyToClipboard: vi.fn() }));
+vi.mock("@/components/ui/use-toast", () => ({ toast: vi.fn() }));
+
+function renderActions(cellId: CellId, name: string) {
+  const store = createStore();
+  const { model, documentAtom, stop } = mockDocumentModel(store);
+  onTestFinished(stop);
+  store.set(requestClientAtom, MockRequestClient.create());
+  store.set(
+    activeNotebookReplicaAtom,
+    partial<NotebookReplica>({
+      documentAtom,
+      notebookAtom: atom(
+        partial<NotebookOutput>({ operations: ["notebook.edit"] }),
+      ),
+    }),
+  );
+  const wrapper = ({ children }: { children: ReactNode }) => (
+    <DocumentWrapper store={store} model={model}>
+      <ModalProvider>{children}</ModalProvider>
+    </DocumentWrapper>
+  );
+  const { result } = renderHook(
+    () =>
+      useCellActionButtons({
+        cell: {
+          cellId,
+          name,
+          config: { disabled: false, hide_code: false, column: null },
+          status: "idle",
+          hasOutput: false,
+          hasConsoleOutput: false,
+          getEditorView: () => null,
+        },
+      }),
+    { wrapper },
+  );
+  return result.current;
+}
+
+function getAction(groups: ActionButton[][], label: string): ActionButton {
+  const action = groups.flat().find((a) => a.label === label);
+  if (!action) {
+    throw new Error(`Action "${label}" not found`);
+  }
+  return action;
+}
+
+describe("useCellActionButtons", () => {
+  beforeEach(() => {
+    vi.clearAllMocks();
+  });
+
+  it("places Copy cell ID right after Copy link to cell", () => {
+    const groups = renderActions(CellId.create(), "_");
+    const linkGroup = groups.find((group) =>
+      group.some((a) => a.label === "Copy link to cell"),
+    );
+    expect(linkGroup?.map((a) => a.label)).toEqual([
+      "Name",
+      "Copy link to cell",
+      "Copy cell ID",
+    ]);
+  });
+
+  it("copies the cell ID of an unnamed cell", async () => {
+    const cellId = CellId.create();
+    const groups = renderActions(cellId, "_");
+
+    await getAction(groups, "Copy cell ID").handle();
+
+    expect(copyToClipboard).toHaveBeenCalledWith(cellId);
+    expect(toast).toHaveBeenCalledWith({
+      description: "Cell ID copied to clipboard",
+    });
+  });
+
+  it("copies the cell ID, not the name, of a named cell", async () => {
+    const cellId = CellId.create();
+    const groups = renderActions(cellId, "my_cell");
+
+    await getAction(groups, "Copy cell ID").handle();
+
+    expect(copyToClipboard).toHaveBeenCalledWith(cellId);
+  });
+
+  it("keeps Copy link to cell disabled for unnamed cells", () => {
+    const groups = renderActions(CellId.create(), "_");
+
+    expect(getAction(groups, "Copy link to cell").disabled).toBe(true);
+    expect(getAction(groups, "Copy cell ID").disabled).toBeUndefined();
+  });
+});

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

@@ -3,13 +3,14 @@
 import React, { useRef, useState } from "react";
 import { Input } from "@/components/ui/input";
 import { Tooltip } from "@/components/ui/tooltip";
-import { getCellNames, useCellActions } from "@/core/cells/cells";
+import { getCellNames } from "@/core/cells/cells";
 import type { CellId } from "@/core/cells/ids";
 import {
   getValidName,
   isInternalCellName,
   normalizeName,
 } from "@/core/cells/names";
+import { useNotebookDocumentEdit } from "@/core/notebook/document";
 import { useOnMount } from "@/hooks/useLifecycle";
 import { cn } from "@/utils/cn";
 import { Events } from "@/utils/events";
@@ -71,9 +72,9 @@ export const NameCellContentEditable: React.FC<{
   value: string;
   className: string;
 }> = ({ value, cellId, className }) => {
-  const { updateCellName } = useCellActions();
+  const edit = useNotebookDocumentEdit();
   const inputProps = useCellNameInput(value, (newName) =>
-    updateCellName({ cellId, name: newName }),
+    edit((tx) => tx.setName(cellId, newName)),
   );
 
   // If the name is the default, don't render the content editable

package/src/components/editor/actions/pair-with-agent-commands.ts

@@ -14,7 +14,7 @@ export const AGENT_LABELS: Record<AgentTab, string> = {
   claude: "Claude",
   codex: "Codex",
   opencode: "OpenCode",
-  prompt: "Prompt",
+  prompt: "Other Agents",
 };
 
 /** How to invoke marimo: from the local checkout in dev, else via uvx. */
@@ -67,7 +67,7 @@ export function getTerminalCommand(
 }
 
 /**
- * The raw prompt for the "Prompt" tab. Mirrors the output of
+ * The raw prompt for the "Other Agents" tab. Mirrors the output of
  * `marimo pair prompt` (see `marimo/_cli/pair/commands.py`) so pasting it into
  * an agent behaves the same as the terminal commands.
  */

package/src/components/editor/actions/pair-with-agent-modal.tsx

@@ -1,7 +1,7 @@
 /* Copyright 2026 Marimo. All rights reserved. */
 
 import { CheckIcon, CopyIcon } from "lucide-react";
-import React, { useEffect, useState } from "react";
+import React, { useEffect, useMemo, useState } from "react";
 import { useAtomValue } from "jotai";
 import { Button } from "@/components/ui/button";
 import {
@@ -53,21 +53,20 @@ export const PairWithAgentModal: React.FC<{
   const config = useAtomValue(pairConfigAtom);
   const hostEnabled = config?.hostEnabled ?? false;
   const stableSessionId = useAtomValue(stableSessionIdAtom);
+  const lookupIdentity = useMemo(
+    () => ({ hostEnabled, stableSessionId }),
+    [hostEnabled, stableSessionId],
+  );
   const [lookup, setLookup] = useState<
     {
-      sessionId: string;
-    } & (
-      | { status: "ready"; notebookId: string | null }
-      | { status: "loading" | "error" }
-    )
+      identity: typeof lookupIdentity;
+    } & ({ status: "ready"; notebookId: string | null } | { status: "error" })
   >();
   useEffect(() => {
-    if (!hostEnabled || !stableSessionId) {
-      setLookup(undefined);
+    const { hostEnabled, stableSessionId: sessionId } = lookupIdentity;
+    if (!hostEnabled || !sessionId) {
       return;
     }
-    const sessionId = stableSessionId;
-    setLookup({ sessionId, status: "loading" });
     let cancelled = false;
     API.get<Record<string, { session_id: string; notebook_id: string | null }>>(
       "/sessions",
@@ -79,22 +78,26 @@ export const PairWithAgentModal: React.FC<{
           );
           setLookup(
             session
-              ? { sessionId, status: "ready", notebookId: session.notebook_id }
-              : { sessionId, status: "error" },
+              ? {
+                  identity: lookupIdentity,
+                  status: "ready",
+                  notebookId: session.notebook_id,
+                }
+              : { identity: lookupIdentity, status: "error" },
           );
         }
       })
       .catch(() => {
         if (!cancelled) {
-          setLookup({ sessionId, status: "error" });
+          setLookup({ identity: lookupIdentity, status: "error" });
         }
       });
     return () => {
       cancelled = true;
     };
-  }, [hostEnabled, stableSessionId]);
+  }, [lookupIdentity]);
   const currentLookup =
-    lookup?.sessionId === stableSessionId ? lookup : undefined;
+    lookup?.identity === lookupIdentity ? lookup : undefined;
   const verified = !hostEnabled || currentLookup?.status === "ready";
   const notebookId =
     hostEnabled && currentLookup?.status === "ready"
@@ -113,7 +116,7 @@ export const PairWithAgentModal: React.FC<{
   return (
     <DialogContent className="sm:max-w-2xl">
       <DialogHeader>
-        <DialogTitle>Pair with an agent</DialogTitle>
+        <DialogTitle>Pair an agent</DialogTitle>
         <DialogDescription>
           Use an AI coding agent to pair-program on this notebook.{" "}
           <a

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

@@ -14,6 +14,7 @@ import {
   DatabaseIcon,
   EyeIcon,
   EyeOffIcon,
+  HashIcon,
   ImageIcon,
   LinkIcon,
   PlayIcon,
@@ -37,8 +38,11 @@ import { Label } from "@/components/ui/label";
 import { toast } from "@/components/ui/use-toast";
 import { aiCompletionCellAtom } from "@/core/ai/state";
 import { maybeAddMarimoImport } from "@/core/cells/add-missing-import";
-import { hasOnlyOneCellAtom, useCellActions } from "@/core/cells/cells";
+import { hasOnlyOneCellAtom } from "@/core/cells/cells";
+import { useCellActions } from "@/core/cells/actions";
+import { showHiddenCodeAtom } from "@/core/cells/editor";
 import { type CellId, SETUP_CELL_ID } from "@/core/cells/ids";
+import { useNotebookDocumentEdit } from "@/core/notebook/document";
 import type { CellData } from "@/core/cells/types";
 import { formatEditorViews } from "@/core/codemirror/format";
 import {
@@ -52,9 +56,8 @@ import {
   appWidthAtom,
   autoInstantiateAtom,
 } from "@/core/config/config";
-import { kioskModeAtom } from "@/core/mode";
-import { useRequestClient } from "@/core/network/requests";
 import type { CellConfig, RuntimeState } from "@/core/network/types";
+import { canEditAtom } from "@/core/notebook";
 import { canLinkToCell, createCellLink } from "@/utils/cell-urls";
 import { copyToClipboard } from "@/utils/copy";
 import { downloadCellOutputAsImage } from "@/utils/download";
@@ -83,16 +86,15 @@ interface Props {
 
 export function useCellActionButtons({ cell, closePopover }: Props) {
   const {
-    createNewCell: createCell,
-    updateCellConfig,
-    updateCellName,
     moveCell,
     sendToTop,
     sendToBottom,
     addColumnBreakpoint,
     clearCellOutput,
-    markUntouched,
+    createCell,
   } = useCellActions();
+  const showHiddenCode = useSetAtom(showHiddenCodeAtom);
+  const edit = useNotebookDocumentEdit();
   const splitCell = useSplitCellCallback();
   const runCell = useRunCell(cell?.cellId);
   const hasOnlyOneCell = useAtomValue(hasOnlyOneCellAtom);
@@ -102,11 +104,10 @@ export function useCellActionButtons({ cell, closePopover }: Props) {
   const setAiCompletionCell = useSetAtom(aiCompletionCellAtom);
   const aiFeaturesEnabled = useAtomValue(aiFeaturesEnabledAtom);
   const autoInstantiate = useAtomValue(autoInstantiateAtom);
-  const kioskMode = useAtomValue(kioskModeAtom);
+  const canEdit = useAtomValue(canEditAtom);
   const appWidth = useAtomValue(appWidthAtom);
-  const { saveCellConfig } = useRequestClient();
 
-  if (!cell || kioskMode) {
+  if (!cell || !canEdit) {
     return [];
   }
 
@@ -120,16 +121,13 @@ export function useCellActionButtons({ cell, closePopover }: Props) {
     status,
   } = cell;
 
-  const toggleDisabled = async () => {
-    const newConfig = { disabled: !config.disabled };
-    await saveCellConfig({ configs: { [cellId]: newConfig } });
-    updateCellConfig({ cellId, config: newConfig });
+  const toggleDisabled = () => {
+    edit((tx) => tx.updateConfig(cellId, { disabled: !config.disabled }));
   };
 
-  const toggleHideCode = async () => {
+  const toggleHideCode = () => {
     const newConfig: Partial<CellConfig> = { hide_code: !config.hide_code };
-    await saveCellConfig({ configs: { [cellId]: newConfig } });
-    updateCellConfig({ cellId, config: newConfig });
+    edit((tx) => tx.updateConfig(cellId, newConfig));
     const editorView = getEditorView();
     // If we're hiding the code, we should blur the editor
     // otherwise, we should focus it
@@ -221,21 +219,19 @@ export function useCellActionButtons({ cell, closePopover }: Props) {
           if (!editorView) {
             return;
           }
-          maybeAddMarimoImport({ autoInstantiate, createNewCell: createCell });
+          maybeAddMarimoImport({ autoInstantiate, createCell });
           switchLanguage(editorView, {
             language: "markdown",
             keepCodeAsIs: false,
           });
           // Code stays visible until the user blurs the cell
           if (!config.hide_code && MARKDOWN_INITIAL_HIDE_CODE) {
-            await saveCellConfig({
-              configs: { [cellId]: { hide_code: MARKDOWN_INITIAL_HIDE_CODE } },
-            });
-            updateCellConfig({
-              cellId,
-              config: { hide_code: MARKDOWN_INITIAL_HIDE_CODE },
-            });
-            markUntouched({ cellId });
+            edit((tx) =>
+              tx.updateConfig(cellId, {
+                hide_code: MARKDOWN_INITIAL_HIDE_CODE,
+              }),
+            );
+            showHiddenCode({ cellId });
           }
         },
         hidden: isSetupCell,
@@ -266,7 +262,7 @@ export function useCellActionButtons({ cell, closePopover }: Props) {
               }
               maybeAddMarimoImport({
                 autoInstantiate,
-                createNewCell: createCell,
+                createCell,
               });
               switchLanguage(editorView, {
                 language: "sql",
@@ -288,7 +284,7 @@ export function useCellActionButtons({ cell, closePopover }: Props) {
         ),
         label: "Create cell above",
         hotkey: "cell.createAbove",
-        handle: () => createCell({ cellId, before: true }),
+        handle: () => createCell({ before: cellId }),
         hidden: isSetupCell,
         redundant: true,
       },
@@ -301,7 +297,7 @@ export function useCellActionButtons({ cell, closePopover }: Props) {
         ),
         label: "Create cell below",
         hotkey: "cell.createBelow",
-        handle: () => createCell({ cellId, before: false }),
+        handle: () => createCell({ after: cellId }),
         redundant: true,
       },
       {
@@ -406,7 +402,7 @@ export function useCellActionButtons({ cell, closePopover }: Props) {
                     }
                   }}
                   onChange={(newName) =>
-                    updateCellName({ cellId, name: newName })
+                    edit((tx) => tx.setName(cellId, newName))
                   }
                 />
               </div>
@@ -417,7 +413,7 @@ export function useCellActionButtons({ cell, closePopover }: Props) {
           <NameCellInput
             placeholder={"cell name"}
             value={name}
-            onChange={(newName) => updateCellName({ cellId, name: newName })}
+            onChange={(newName) => edit((tx) => tx.setName(cellId, newName))}
             onEnterKey={() => closePopover?.()}
           />
         ),
@@ -435,6 +431,15 @@ export function useCellActionButtons({ cell, closePopover }: Props) {
           toast({ description: "Link copied to clipboard" });
         },
       },
+      {
+        icon: <HashIcon size={13} strokeWidth={1.5} />,
+        label: "Copy cell ID",
+        additionalKeywords: ["cell_id", "identifier"],
+        handle: async () => {
+          await copyToClipboard(cellId);
+          toast({ description: "Cell ID copied to clipboard" });
+        },
+      },
     ],
 
     // Delete

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],
   );
 };