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

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

@@ -3,18 +3,23 @@ import { historyField } from "@codemirror/commands";
 import { EditorState, type Extension, StateEffect } from "@codemirror/state";
 import { EditorView, ViewPlugin } from "@codemirror/view";
 import { useIntersectionObserver } from "@uidotdev/usehooks";
-import { useAtom, useAtomValue } from "jotai";
+import { useAtom, useAtomValue, useSetAtom } from "jotai";
 import React, { memo, useEffect, useMemo, useRef, useState } from "react";
 import useEvent from "react-use-event-hook";
-import { Button } from "@/components/ui/button";
-import { DelayMount } from "@/components/utils/delay-mount";
 import { aiCompletionCellAtom } from "@/core/ai/state";
 import { maybeAddMarimoImport } from "@/core/cells/add-missing-import";
-import { getNotebook, useCellActions } from "@/core/cells/cells";
-import { SETUP_CELL_ID } from "@/core/cells/ids";
+import { useCellActions } from "@/core/cells/actions";
+import {
+  showHiddenCodeAtom,
+  takeRestoredEditorStateAtom,
+  useRestoredEditorState,
+} from "@/core/cells/editor";
+import {
+  useNotebookCellCodeExtension,
+  useNotebookDocumentEdit,
+} from "@/core/notebook/document";
 import { usePendingDeleteService } from "@/core/cells/pending-delete-service";
 import type { CellData, CellRuntimeState } from "@/core/cells/types";
-import { notebookCellEditorViews } from "@/core/cells/utils";
 import { setupCodeMirror } from "@/core/codemirror/cm";
 import { acceptCompletionOnEnterAtom } from "@/core/codemirror/completion/accept-on-enter-atom";
 import { editorMountScheduler } from "@/core/codemirror/editor-mount-scheduler";
@@ -24,24 +29,12 @@ import {
   reconfigureLanguageEffect,
   switchLanguage,
 } from "@/core/codemirror/language/extension";
-import {
-  getEditorCodeAsPython,
-  updateEditorCodeFromPython,
-} from "@/core/codemirror/language/utils";
 import { MARKDOWN_INITIAL_HIDE_CODE } from "@/core/codemirror/language/languages/markdown";
 import type { LanguageAdapterType } from "@/core/codemirror/language/types";
-import {
-  connectedDocAtom,
-  realTimeCollaboration,
-} from "@/core/codemirror/rtc/extension";
 import { autoInstantiateAtom, isAiFeatureEnabled } from "@/core/config/config";
 import type { UserConfig } from "@/core/config/config-schema";
 import { OverridingHotkeyProvider } from "@/core/hotkeys/hotkeys";
-import { connectionAtom } from "@/core/network/connection";
-import { useRequestClient } from "@/core/network/requests";
-import { canUseRtc, isRtcEnabled } from "@/core/rtc/state";
 import { useSaveNotebook } from "@/core/saving/save-component";
-import { isAppConnecting } from "@/core/websocket/connection-utils";
 import type { Theme } from "@/theme/useTheme";
 import { cn } from "@/utils/cn";
 import { invariant } from "@/utils/invariant";
@@ -59,7 +52,7 @@ import { LanguageToggles } from "./language-toggle";
 export interface CellEditorProps
   extends
     Pick<CellRuntimeState, "status">,
-    Pick<CellData, "id" | "code" | "serializedEditorState" | "config"> {
+    Pick<CellData, "id" | "code" | "config"> {
   runCell: () => void;
   theme: Theme;
   showPlaceholder: boolean;
@@ -93,21 +86,8 @@ export interface CellEditorProps
    * defaults to `#App`.
    */
   tooltipParentSelector?: string;
-}
-
-export function composeRtcEditorConfig(
-  code: string,
-  extensions: Extension[],
-  rtc: ReturnType<typeof realTimeCollaboration> | undefined,
-): { code: string; extensions: Extension[] } {
-  if (!rtc) {
-    return { code, extensions };
-  }
-
-  return {
-    code: rtc.code,
-    extensions: [...extensions, rtc.extension],
-  };
+  /** Further editor extensions, for an editor outside the notebook. */
+  extensions?: Extension[];
 }
 
 const CellEditorInternal = ({
@@ -117,7 +97,6 @@ const CellEditorInternal = ({
   config: cellConfig,
   code,
   status,
-  serializedEditorState,
   setEditorView,
   runCell,
   userConfig,
@@ -132,12 +111,17 @@ const CellEditorInternal = ({
   inlineAiTooltip,
   outputArea,
   tooltipParentSelector,
+  extensions: extraExtensions,
 }: CellEditorProps) => {
   const [aiCompletionCell, setAiCompletionCell] = useAtom(aiCompletionCellAtom);
+  const serializedEditorState = useRestoredEditorState(cellId);
+  const takeRestoredEditorState = useSetAtom(takeRestoredEditorStateAtom);
+  const showHiddenCodeNow = useSetAtom(showHiddenCodeAtom);
   const deleteCell = useDeleteCellCallback();
   const { saveOrNameNotebook } = useSaveNotebook();
   const pendingDeleteService = usePendingDeleteService();
-  const { saveCellConfig } = useRequestClient();
+  const edit = useNotebookDocumentEdit();
+  const codeBinding = useNotebookCellCodeExtension(cellId);
 
   const loading = status === "running" || status === "queued";
   const cellActions = useCellActions();
@@ -177,9 +161,7 @@ const CellEditorInternal = ({
   const toggleHideCode = useEvent(() => {
     // Use cellConfig.hide_code instead of hidden, since it may be temporarily shown
     const nextHidden = !cellConfig.hide_code;
-    // Fire-and-forget save
-    void saveCellConfig({ configs: { [cellId]: { hide_code: nextHidden } } });
-    cellActions.updateCellConfig({ cellId, config: { hide_code: nextHidden } });
+    edit((tx) => tx.updateConfig(cellId, { hide_code: nextHidden }));
     return nextHidden;
   });
 
@@ -188,25 +170,21 @@ const CellEditorInternal = ({
   const afterToggleMarkdown = useEvent(() => {
     maybeAddMarimoImport({
       autoInstantiate,
-      createNewCell: cellActions.createNewCell,
+      createCell: cellActions.createCell,
     });
     // Code stays visible until the user blurs the cell
     if (!cellConfig.hide_code && MARKDOWN_INITIAL_HIDE_CODE) {
-      void saveCellConfig({
-        configs: { [cellId]: { hide_code: MARKDOWN_INITIAL_HIDE_CODE } },
-      });
-      cellActions.updateCellConfig({
-        cellId,
-        config: { hide_code: MARKDOWN_INITIAL_HIDE_CODE },
-      });
-      cellActions.markUntouched({ cellId });
+      edit((tx) =>
+        tx.updateConfig(cellId, { hide_code: MARKDOWN_INITIAL_HIDE_CODE }),
+      );
+      showHiddenCodeNow({ cellId });
     }
   });
 
   const afterToggleSQL = useEvent(() => {
     maybeAddMarimoImport({
       autoInstantiate,
-      createNewCell: cellActions.createNewCell,
+      createCell: cellActions.createCell,
     });
   });
 
@@ -226,41 +204,13 @@ const CellEditorInternal = ({
         saveNotebook: saveOrNameNotebook,
         createManyBelow: (cells) => {
           for (const code of cells.toReversed()) {
-            cellActions.createNewCell({
+            cellActions.createCell({
+              after: cellId,
               code,
-              before: false,
-              cellId: cellId,
-              // If the code already exists, skip creation
               skipIfCodeExists: true,
             });
           }
         },
-        addOrAppendSetupCell: (code) => {
-          const notebook = getNotebook();
-          // No setup cell yet: create one with this code at the top.
-          if (!notebook.cellIds.setupCellExists()) {
-            cellActions.addSetupCellIfDoesntExist({ code });
-            return;
-          }
-          // Setup cell exists: append the pasted code to it. Prefer updating
-          // the mounted editor (keeps the view and state in sync, like
-          // formatting does); fall back to a plain state update otherwise.
-          const view = notebookCellEditorViews(notebook)[SETUP_CELL_ID];
-          const existing = view
-            ? getEditorCodeAsPython(view)
-            : notebook.cellData[SETUP_CELL_ID].code;
-          const merged = existing.trim()
-            ? `${existing.trimEnd()}\n\n${code}`
-            : code;
-          if (view) {
-            updateEditorCodeFromPython(view, merged);
-          }
-          cellActions.updateCellCode({
-            cellId: SETUP_CELL_ID,
-            code: merged,
-            formattingChange: false,
-          });
-        },
         splitCell,
         toggleHideCode,
         aiCellCompletion: () => {
@@ -293,6 +243,8 @@ const CellEditorInternal = ({
     });
 
     extensions.push(
+      codeBinding,
+      ...(extraExtensions ?? []),
       // Listen to code changes if we can use markdown
       // Also update the language adapter
       ViewPlugin.define((view) => {
@@ -332,6 +284,8 @@ const CellEditorInternal = ({
     return extensions;
   }, [
     cellId,
+    codeBinding,
+    extraExtensions,
     acceptCompletionOnEnter,
     userConfig.keymap,
     userConfig.completion,
@@ -357,50 +311,24 @@ const CellEditorInternal = ({
     saveOrNameNotebook,
   ]);
 
-  const rtcEnabled = isRtcEnabled() && canUseRtc(cellId);
-  const createEditorConfig = (initialCode?: string) => {
-    if (!rtcEnabled) {
-      return composeRtcEditorConfig(code, extensions, undefined);
-    }
-
-    const rtc = realTimeCollaboration(
-      cellId,
-      (code) => {
-        // It's not really a formatting change,
-        // but this means it won't be marked as stale
-        cellActions.updateCellCode({ cellId, code, formattingChange: true });
-      },
-      initialCode,
-    );
-    return composeRtcEditorConfig(code, extensions, rtc);
-  };
-
   const handleInitializeEditor = useEvent(() => {
-    const editorConfig = createEditorConfig(code);
-
     // Create a new editor
     const ev = new EditorView({
-      state: EditorState.create({
-        doc: editorConfig.code,
-        extensions: editorConfig.extensions,
-      }),
+      state: EditorState.create({ doc: code, extensions }),
     });
     setEditorView(ev);
     // Initialize the language adapter
-    if (!rtcEnabled) {
-      switchLanguage(ev, {
-        language: getInitialLanguageAdapter(ev.state).type,
-      });
-    }
+    switchLanguage(ev, {
+      language: getInitialLanguageAdapter(ev.state).type,
+    });
   });
 
   const handleReconfigureEditor = useEvent(() => {
     invariant(editorViewRef.current !== null, "Editor view is not initialized");
-    const editorConfig = createEditorConfig();
 
     editorViewRef.current.dispatch({
       effects: [
-        StateEffect.reconfigure.of([editorConfig.extensions]),
+        StateEffect.reconfigure.of([extensions]),
         reconfigureLanguageEffect(editorViewRef.current, {
           completionConfig: userConfig.completion,
           hotkeysProvider: new OverridingHotkeyProvider(
@@ -417,27 +345,21 @@ const CellEditorInternal = ({
 
   const handleDeserializeEditor = useEvent(() => {
     invariant(serializedEditorState, "Editor view is not initialized");
-    const editorConfig = createEditorConfig(code);
 
     const ev = new EditorView({
       state: EditorState.fromJSON(
         serializedEditorState,
-        {
-          doc: editorConfig.code,
-          extensions: editorConfig.extensions,
-        },
+        { doc: code, extensions },
         { history: historyField },
       ),
     });
     // Initialize the language adapter
-    if (!rtcEnabled) {
-      switchLanguage(ev, {
-        language: getInitialLanguageAdapter(ev.state).type,
-      });
-    }
+    switchLanguage(ev, {
+      language: getInitialLanguageAdapter(ev.state).type,
+    });
     setEditorView(ev);
-    // Clear the serialized state so that we don't re-create the editor next time
-    cellActions.clearSerializedEditorState({ cellId });
+    // The restored state is used once; the next mount builds a fresh editor.
+    takeRestoredEditorState({ cellId });
   });
   const [isEditorMounted, setIsEditorMounted] = useState(
     serializedEditorState !== null,
@@ -572,7 +494,7 @@ const CellEditorInternal = ({
           <div className="absolute top-1 right-5">
             <LanguageToggles
               code={code}
-              editorView={editorViewRef.current}
+              editorViewRef={editorViewRef}
               currentLanguageAdapter={languageAdapter}
               onAfterToggle={afterToggleMarkdown}
             />
@@ -633,45 +555,4 @@ const CellCodeMirrorEditor = React.forwardRef(
 );
 CellCodeMirrorEditor.displayName = "CellCodeMirrorEditor";
 
-// Wait until the websocket connection is open before rendering the editor
-// This is used for real-time collaboration since the backend needs the connection started
-// before connecting the rtc websockets
-function WithWaitUntilConnected<T extends Pick<CellEditorProps, "id">>(
-  Component: React.ComponentType<T>,
-) {
-  const WaitUntilConnectedComponent = (props: T) => {
-    const connection = useAtomValue(connectionAtom);
-    const [rtcDoc, setRtcDoc] = useAtom(connectedDocAtom);
-
-    if (
-      isRtcEnabled() &&
-      canUseRtc(props.id) &&
-      (isAppConnecting(connection.state) || rtcDoc === undefined)
-    ) {
-      return (
-        <div className="flex h-full w-full items-baseline p-4">
-          <DelayMount milliseconds={1000} fallback={null}>
-            <span>Waiting for real-time collaboration connection...</span>
-            <Button
-              variant="link"
-              onClick={() => {
-                setRtcDoc("disabled");
-              }}
-            >
-              Turn off real-time collaboration
-            </Button>
-          </DelayMount>
-        </div>
-      );
-    }
-
-    return <Component {...props} />;
-  };
-
-  WaitUntilConnectedComponent.displayName = `WithWaitUntilConnected(${Component.displayName})`;
-  return WaitUntilConnectedComponent;
-}
-
-export const CellEditor = isRtcEnabled()
-  ? WithWaitUntilConnected(memo(CellEditorInternal))
-  : memo(CellEditorInternal);
+export const CellEditor = memo(CellEditorInternal);

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

@@ -13,7 +13,7 @@ import { Functions } from "@/utils/functions";
 import { MarkdownIcon, PythonIcon } from "./icons";
 
 interface LanguageTogglesProps {
-  editorView: EditorView | null;
+  editorViewRef: React.RefObject<EditorView | null>;
   code: string;
   currentLanguageAdapter: LanguageAdapter["type"] | undefined;
   onAfterToggle: () => void;
@@ -25,7 +25,7 @@ interface LanguageTogglesProps {
 }
 
 export const LanguageToggles: React.FC<LanguageTogglesProps> = ({
-  editorView,
+  editorViewRef,
   code,
   currentLanguageAdapter,
   onAfterToggle,
@@ -43,7 +43,7 @@ export const LanguageToggles: React.FC<LanguageTogglesProps> = ({
   return (
     <div className={className}>
       <LanguageToggle
-        editorView={editorView}
+        editorViewRef={editorViewRef}
         currentLanguageAdapter={currentLanguageAdapter}
         canSwitchToLanguage={canUseSQL && currentLanguageAdapter === "python"}
         icon={
@@ -58,7 +58,7 @@ export const LanguageToggles: React.FC<LanguageTogglesProps> = ({
         onAfterToggle={onAfterToggle}
       />
       <LanguageToggle
-        editorView={editorView}
+        editorViewRef={editorViewRef}
         currentLanguageAdapter={currentLanguageAdapter}
         canSwitchToLanguage={
           canUseMarkdown && currentLanguageAdapter === "python"
@@ -75,7 +75,7 @@ export const LanguageToggles: React.FC<LanguageTogglesProps> = ({
         onAfterToggle={onAfterToggle}
       />
       <LanguageToggle
-        editorView={editorView}
+        editorViewRef={editorViewRef}
         currentLanguageAdapter={currentLanguageAdapter}
         canSwitchToLanguage={true}
         icon={
@@ -95,7 +95,7 @@ export const LanguageToggles: React.FC<LanguageTogglesProps> = ({
 
 interface Props {
   className?: string;
-  editorView: EditorView | null;
+  editorViewRef: React.RefObject<EditorView | null>;
   canSwitchToLanguage: boolean;
   currentLanguageAdapter: LanguageAdapter["type"] | undefined;
   toType: LanguageAdapter["type"];
@@ -105,7 +105,7 @@ interface Props {
 }
 
 export const LanguageToggle: React.FC<Props> = ({
-  editorView,
+  editorViewRef,
   currentLanguageAdapter,
   canSwitchToLanguage,
   icon,
@@ -114,6 +114,7 @@ export const LanguageToggle: React.FC<Props> = ({
   onAfterToggle,
 }) => {
   const handleClick = () => {
+    const editorView = editorViewRef.current;
     if (!editorView) {
       return;
     }

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

@@ -14,6 +14,7 @@ import { Toolbar } from "@/components/layout/toolbar";
 import { Button } from "@/components/ui/button";
 import { Tooltip } from "@/components/ui/tooltip";
 import { useNotebook } from "@/core/cells/cells";
+import { columnOf, useLayout } from "@/core/cells/layout";
 import type { CellId } from "@/core/cells/ids";
 import { getDescendantsStatus } from "@/core/cells/utils";
 import { cn } from "@/utils/cn";
@@ -64,13 +65,19 @@ export const CollapsedCellBanner: React.FC<{
   count: number;
 }> = memo(({ onClick, count, cellId }) => {
   const notebook = useNotebook();
-  const states = getDescendantsStatus(notebook, cellId);
+  const layout = useLayout();
+  const states = getDescendantsStatus(
+    notebook,
+    columnOf(layout, cellId).folds.get(cellId) ?? [],
+  );
 
   return (
-    <div
+    <button
+      type="button"
+      aria-label={`Expand ${count} collapsed ${count === 1 ? "cell" : "cells"}`}
       onClick={onClick}
       className={cn(
-        "flex items-center justify-between w-[calc(100%-2rem)] h-9 bg-muted rounded-b mx-4 opacity-80 hover:opacity-100 cursor-pointer",
+        "flex items-center justify-between w-[calc(100%-2rem)] h-9 bg-muted rounded-b mx-4 opacity-80 hover:opacity-100 cursor-pointer focus-visible:outline focus-visible:outline-ring",
       )}
     >
       <Toolbar
@@ -102,7 +109,7 @@ export const CollapsedCellBanner: React.FC<{
           </>
         }
       />
-    </div>
+    </button>
   );
 });
 CollapsedCellBanner.displayName = "CollapsedCellBanner";

package/src/components/editor/cell/useAddCell.ts

@@ -2,28 +2,24 @@
 
 import { useAtomValue } from "jotai";
 import { maybeAddAltairImport } from "@/core/cells/add-missing-import";
-import { useCellActions } from "@/core/cells/cells";
+import { useCellActions } from "@/core/cells/actions";
 import { useLastFocusedCellId } from "@/core/cells/focus";
 import { autoInstantiateAtom } from "@/core/config/config";
 
 export function useAddCodeToNewCell(): (code: string) => void {
   const autoInstantiate = useAtomValue(autoInstantiateAtom);
   const lastFocusedCellId = useLastFocusedCellId();
-  const { createNewCell } = useCellActions();
+  const { createCell } = useCellActions();
 
   return (code: string) => {
     if (code.includes("alt")) {
       maybeAddAltairImport({
         autoInstantiate,
-        createNewCell,
+        createCell,
         fromCellId: lastFocusedCellId,
       });
     }
 
-    createNewCell({
-      code: code,
-      before: false,
-      cellId: lastFocusedCellId ?? "__end__",
-    });
+    createCell({ after: lastFocusedCellId, code });
   };
 }

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

@@ -3,35 +3,22 @@
 import useEvent from "react-use-event-hook";
 import { UndoButton } from "@/components/buttons/undo-button";
 import { toast } from "@/components/ui/use-toast";
-import {
-  hasOnlyOneCellAtom,
-  notebookAtom,
-  useCellActions,
-} from "@/core/cells/cells";
+import { useCellActions } from "@/core/cells/actions";
+import { notebookAtom } from "@/core/cells/cells";
 import type { CellId } from "@/core/cells/ids";
-import { useRequestClient } from "@/core/network/requests";
 import { store } from "@/core/state/jotai";
 
 export function useDeleteCellCallback() {
-  const { deleteCell, undoDeleteCell } = useCellActions();
-  const { sendDeleteCell } = useRequestClient();
+  const { deleteCells, undoDeleteCell } = useCellActions();
 
   return useEvent((opts: { cellId: CellId }) => {
-    // Can't delete the last cell
-    if (store.get(hasOnlyOneCellAtom)) {
-      return;
-    }
-
     const { cellId } = opts;
     const notebook = store.get(notebookAtom);
     const isEmptyCell = (notebook.cellData[cellId]?.code ?? "").trim() === "";
 
-    // Optimistic update
-    deleteCell({ cellId });
-    sendDeleteCell({ cellId: cellId }).catch(() => {
-      // Fall back on failure
-      undoDeleteCell();
-    });
+    if (!deleteCells({ cellIds: [cellId] })) {
+      return;
+    }
 
     if (!isEmptyCell) {
       const { dismiss } = toast({
@@ -53,33 +40,20 @@ export function useDeleteCellCallback() {
 }
 
 export function useDeleteManyCellsCallback() {
-  const { deleteCell, undoDeleteCell } = useCellActions();
-  const { sendDeleteCell } = useRequestClient();
+  const { deleteCells, undoDeleteCell } = useCellActions();
 
-  return useEvent(async (opts: { cellIds: CellId[] }) => {
-    // Can't delete the last cell
-    if (store.get(hasOnlyOneCellAtom)) {
+  return useEvent((opts: { cellIds: CellId[] }) => {
+    if (!deleteCells({ cellIds: opts.cellIds })) {
       return;
     }
 
-    const { cellIds } = opts;
-    for (const cellId of cellIds) {
-      await sendDeleteCell({ cellId }).then(() => {
-        deleteCell({ cellId });
-      });
-    }
-
     const { dismiss } = toast({
       title: "Cells deleted",
       action: (
         <UndoButton
           data-testid="undo-delete-button"
           onClick={() => {
-            for (const _cellId of cellIds) {
-              // This function does not take a cellId,
-              // so we just undo the number of cells that were deleted
-              undoDeleteCell();
-            }
+            undoDeleteCell();
             dismiss();
           }}
         />

package/src/components/editor/cell/useRunCells.ts

@@ -3,11 +3,8 @@
 import { closeCompletion } from "@codemirror/autocomplete";
 import { atom, useSetAtom } from "jotai";
 import useEvent from "react-use-event-hook";
-import {
-  getNotebook,
-  type NotebookState,
-  useCellActions,
-} from "@/core/cells/cells";
+import { getNotebook, type NotebookState } from "@/core/cells/cells";
+import { useCellActions } from "@/core/cells/actions";
 import type { CellId } from "@/core/cells/ids";
 import { enabledCellIds, staleCellIds } from "@/core/cells/utils";
 import { getCurrentLanguageAdapter } from "@/core/codemirror/language/commands";

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

@@ -1,5 +1,5 @@
 /* Copyright 2026 Marimo. All rights reserved. */
-import { useEffect, useState } from "react";
+import { useDelayElapsed } from "@/hooks/useDelayElapsed";
 
 const DELAY_MS = 200;
 
@@ -10,17 +10,7 @@ const DELAY_MS = 200;
 export function useShouldShowInterrupt(running: boolean) {
   // Start a timer when the run starts.
   // After 200ms, show the interrupt button to avoid flickering.
-  const [hasRunLongEnough, setHasRunLongEnough] = useState(false);
-  useEffect(() => {
-    if (!running) {
-      return;
-    }
-    setHasRunLongEnough(false);
-    const timeout = setTimeout(() => {
-      setHasRunLongEnough(true);
-    }, DELAY_MS);
-    return () => clearTimeout(timeout);
-  }, [running]);
+  const hasRunLongEnough = useDelayElapsed(running ? DELAY_MS : 0);
 
   return running && hasRunLongEnough;
 }

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

@@ -3,23 +3,26 @@
 import useEvent from "react-use-event-hook";
 import { UndoButton } from "@/components/buttons/undo-button";
 import { toast } from "@/components/ui/use-toast";
-import { getCellEditorView, useCellActions } from "@/core/cells/cells";
+import { getCellEditorView } from "@/core/cells/cells";
+import { useCellActions } from "@/core/cells/actions";
 import type { CellId } from "@/core/cells/ids";
 import { getEditorCodeAsPython } from "@/core/codemirror/language/utils";
 
 export function useSplitCellCallback() {
   const { splitCell, undoSplitCell } = useCellActions();
 
-  return useEvent((opts: { cellId: CellId }) => {
+  return useEvent((opts: { cellId: CellId }): CellId | null => {
     // Save snapshot of code for undo
     const cellEditorView = getCellEditorView(opts.cellId);
     if (!cellEditorView) {
-      return;
+      return null;
     }
     const code = getEditorCodeAsPython(cellEditorView);
 
-    // Optimistic update
-    splitCell(opts);
+    const newCellId = splitCell(opts);
+    if (newCellId === null) {
+      return null;
+    }
 
     const { dismiss } = toast({
       title: "Cell split",
@@ -29,7 +32,7 @@ export function useSplitCellCallback() {
           size="sm"
           variant="outline"
           onClick={() => {
-            undoSplitCell({ ...opts, snapshot: code });
+            undoSplitCell({ ...opts, newCellId, snapshot: code });
             dismiss();
           }}
         >
@@ -37,5 +40,6 @@ export function useSplitCellCallback() {
         </UndoButton>
       ),
     });
+    return newCellId;
   });
 }