@marimo-team/islands 0.25.3-dev24 → 0.25.3-dev244

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}

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();
+    }
+  });
+});

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);