@marimo-team/islands 0.25.3-dev170 → 0.25.3-dev175

package/src/core/cells/document.ts

@@ -10,11 +10,13 @@
 
 import type * as Protocol from "@marimo-team/client/protocol";
 import { dequal as isEqual } from "dequal";
+import type { createStore } from "jotai";
 import { createRef, type RefObject } from "react";
 import type { CellHandle } from "@/components/editor/notebook-cell";
 import type {
   DocumentEdit,
   DocumentEditor,
+  DocumentModel,
   EditableCell,
   EditableDocument,
 } from "@/core/notebook/document";
@@ -25,7 +27,7 @@ import type { CellConfig } from "../network/types";
 import { canUseRtc, isRtcEnabled } from "../rtc/state";
 import { store } from "../state/jotai";
 import { activeDocumentModelAtom } from "../session/active-notebook";
-import type { CellActions, NotebookState } from "./cells";
+import type { NotebookState, ReducerCellActions } from "./cells";
 import { groupByColumn } from "./columns";
 import type { CellId } from "./ids";
 import {
@@ -36,7 +38,7 @@ import {
   isEdited,
 } from "./types";
 
-type CellAction = DispatchedActionOf<CellActions>;
+type CellAction = DispatchedActionOf<ReducerCellActions>;
 
 // Following the document ----------------------------------------------------
 
@@ -124,14 +126,17 @@ export function applyDocument(
       continue;
     }
     if (current === undefined) {
-      // A new cell has no run to be edited from.
+      // A new cell with code is stale until it runs. The kernel's last run
+      // settles it on kernel-ready.
+      const code = cell.code ?? "";
       cellData = {
         ...cellData,
         [id]: createCell({
           id,
-          code: cell.code ?? "",
+          code,
           name: cell.name,
           config: createCellConfig(cell.config),
+          edited: code.trim().length > 0,
         }),
       };
       cellRuntime = { ...cellRuntime, [id]: createCellRuntimeState() };
@@ -165,6 +170,33 @@ export function applyDocument(
   return { ...state, cellIds, cellData, cellRuntime, cellHandles };
 }
 
+/**
+ * Reconcile the model's document now and on each change, until the signal aborts.
+ *
+ * TODO(MO-8028): Delete with `applyDocument`.
+ */
+export function followDocument({
+  store,
+  model,
+  reconcile,
+  signal,
+}: {
+  store: ReturnType<typeof createStore>;
+  model: Pick<DocumentModel, "documentAtom">;
+  reconcile: (document: EditableDocument) => void;
+  signal: AbortSignal;
+}): void {
+  const follow = () => {
+    const document = store.get(model.documentAtom);
+    if (document !== null) {
+      reconcile(document);
+    }
+  };
+  const unsubscribe = store.sub(model.documentAtom, follow);
+  signal.addEventListener("abort", unsubscribe, { once: true });
+  follow();
+}
+
 // Editing the document ------------------------------------------------------
 
 const CONFIG_FIELDS = ["disabled", "hide_code", "expand_output"] as const;
@@ -291,8 +323,8 @@ export function documentEditMiddleware(
   next: NotebookState,
   action: CellAction,
 ): void {
-  // The kernel's cells replace the notebook rather than edit it.
-  if (action.type === "setCells") {
+  // The document replaces the notebook's cells, so following it is no edit.
+  if (action.type === "reconcileDocument") {
     return;
   }
   if (prev.cellIds === next.cellIds && prev.cellData === next.cellData) {

package/src/core/cells/outputs.ts

@@ -1,15 +1,15 @@
 /* Copyright 2026 Marimo. All rights reserved. */
 
 import { useAtomValue } from "jotai";
-import { selectAtom } from "jotai/utils";
-import { useEffect, useMemo, useState } from "react";
+import { useEffect, useState } from "react";
 import { useEvent } from "@/hooks/useEvent";
-import { Logger } from "@/utils/Logger";
+import {
+  useNotebookCellConfig,
+  useNotebookDocumentEdit,
+} from "../notebook/document";
 import type { OutputMessage } from "../kernel/messages";
 import { viewStateAtom } from "../mode";
-import { getRequestClient } from "../network/requests";
-import { cellDataAtom, useCellActions } from "./cells";
-import type { CellId } from "./ids";
+import { type CellId, SCRATCH_CELL_ID } from "./ids";
 
 // This does not need to be overcomplicated. We can just store the expanded
 // state in a global map instead of Jotai since state is not shared between cells.
@@ -27,20 +27,9 @@ const expandedConsoleOutputs: Record<CellId, boolean> = {};
  */
 export function useExpandedOutput(cellId: CellId) {
   const { mode } = useAtomValue(viewStateAtom);
-  const { updateCellConfig } = useCellActions();
-  // Select just the flag: every output would otherwise re-render on each
-  // keystroke in its cell.
-  const isConfigExpanded = useAtomValue(
-    useMemo(
-      () =>
-        selectAtom(
-          cellDataAtom(cellId),
-          (data) => data?.config.expand_output ?? false,
-          Object.is,
-        ),
-      [cellId],
-    ),
-  );
+  const edit = useNotebookDocumentEdit();
+  const isConfigExpanded =
+    useNotebookCellConfig(cellId)?.expand_output ?? false;
   const [isSessionExpanded, setIsSessionExpanded] = useState(
     () => expandedOutputs[cellId] ?? isConfigExpanded,
   );
@@ -55,16 +44,11 @@ export function useExpandedOutput(cellId: CellId) {
 
   const setIsExpanded = useEvent((expanded: boolean) => {
     setIsSessionExpanded(expanded);
-    if (!isEditable) {
+    // The scratchpad is not in the document.
+    if (!isEditable || cellId === SCRATCH_CELL_ID) {
       return;
     }
-    const config = { expand_output: expanded };
-    updateCellConfig({ cellId, config });
-    getRequestClient()
-      .saveCellConfig({ configs: { [cellId]: config } })
-      .catch((error) => {
-        Logger.error("Failed to save expand_output config", error);
-      });
+    edit((tx) => tx.updateConfig(cellId, { expand_output: expanded }));
   });
 
   return [isExpanded, setIsExpanded] as const;

package/src/core/codemirror/cells/extensions.ts

@@ -190,7 +190,7 @@ function cellKeymaps({
 
           if (isAtStartOfEditor(ev)) {
             const actions = ev.state.facet(cellActionsState);
-            actions.moveToNextCell({ cellId, before: true, noCreate: true });
+            actions.moveToNextCell({ cellId, before: true, create: false });
             return true;
           }
           return false;
@@ -207,7 +207,7 @@ function cellKeymaps({
 
           if (isAtEndOfEditor(ev)) {
             const actions = ev.state.facet(cellActionsState);
-            actions.moveToNextCell({ cellId, before: false, noCreate: true });
+            actions.moveToNextCell({ cellId, before: false, create: false });
             return true;
           }
           return false;
@@ -219,7 +219,7 @@ function cellKeymaps({
         stopPropagation: true,
         run: (ev) => {
           const actions = ev.state.facet(cellActionsState);
-          actions.moveToNextCell({ cellId, before: false, noCreate: true });
+          actions.moveToNextCell({ cellId, before: false, create: false });
           return true;
         },
       },
@@ -229,7 +229,7 @@ function cellKeymaps({
         stopPropagation: true,
         run: (ev) => {
           const actions = ev.state.facet(cellActionsState);
-          actions.moveToNextCell({ cellId, before: true, noCreate: true });
+          actions.moveToNextCell({ cellId, before: true, create: false });
           return true;
         },
       },
@@ -260,7 +260,7 @@ function cellKeymaps({
         run: (ev) => {
           ev.contentDOM.blur();
           const actions = ev.state.facet(cellActionsState);
-          actions.createNewCell({ cellId, before: true });
+          actions.createCell({ before: cellId });
           return true;
         },
       },
@@ -271,7 +271,7 @@ function cellKeymaps({
         run: (ev) => {
           ev.contentDOM.blur();
           const actions = ev.state.facet(cellActionsState);
-          actions.createNewCell({ cellId, before: false });
+          actions.createCell({ after: cellId });
           return true;
         },
       },

package/src/core/codemirror/cells/state.ts

@@ -1,7 +1,7 @@
 /* Copyright 2026 Marimo. All rights reserved. */
 
 import { Facet } from "@codemirror/state";
-import type { CellActions } from "@/core/cells/cells";
+import type { CellActions } from "@/core/cells/actions";
 import type { CellId } from "@/core/cells/ids";
 
 export interface CodemirrorCellActions extends CellActions {

package/src/core/codemirror/keymaps/vim.ts

@@ -37,7 +37,7 @@ export function vimKeymapExtension(): Extension[] {
           if (isAtEndOfEditor(ev, true) && isInVimNormalMode(ev)) {
             const actions = ev.state.facet(cellActionsState);
             const cellId = ev.state.facet(cellIdState);
-            actions.moveToNextCell({ cellId, before: false, noCreate: true });
+            actions.moveToNextCell({ cellId, before: false, create: false });
             return true;
           }
           return false;
@@ -51,7 +51,7 @@ export function vimKeymapExtension(): Extension[] {
           if (isAtStartOfEditor(ev) && isInVimNormalMode(ev)) {
             const actions = ev.state.facet(cellActionsState);
             const cellId = ev.state.facet(cellIdState);
-            actions.moveToNextCell({ cellId, before: true, noCreate: true });
+            actions.moveToNextCell({ cellId, before: true, create: false });
             return true;
           }
           return false;

package/src/core/edit-app.tsx

@@ -5,8 +5,8 @@ import { Tooltip } from "radix-ui";
 
 const TooltipProvider = Tooltip.Provider;
 
-import { useAtomValue, useSetAtom } from "jotai";
-import { useEffect } from "react";
+import { useAtomValue } from "jotai";
+import { lazy, useEffect } from "react";
 import { NotStartedConnectionAlert } from "@/components/editor/alerts/connecting-alert";
 import { Controls } from "@/components/editor/controls/Controls";
 import { SandboxController } from "@/components/editor/chrome/panels/sandbox-controller";
@@ -31,20 +31,21 @@ import {
   hasCellsAtom,
   notebookIsRunningAtom,
   numColumnsAtom,
-  useCellActions,
 } from "./cells/cells";
+import { useCellActions } from "./cells/actions";
 import type { AppConfig, UserConfig } from "./config/config-schema";
-import { DocumentProvider } from "./notebook/document";
 import { RuntimeState } from "./kernel/RuntimeState";
-import { getSessionId } from "./kernel/session";
 import { useTogglePresenting } from "./layout/useTogglePresenting";
 import { viewStateAtom } from "./mode";
 import { useRequestClient } from "./network/requests";
 import { useFilename } from "./saving/filename";
-import { lastSavedNotebookAtom } from "./saving/state";
-import { useNotebookSession } from "./session/useNotebookSession";
+import type { ConnectionStatus } from "./session/types";
 
-interface AppProps {
+const LazyCommandPalette = lazy(
+  () => import("../components/editor/controls/command-palette"),
+);
+
+export interface EditAppProps {
   /**
    * The user config.
    */
@@ -57,19 +58,24 @@ interface AppProps {
    * If true, the floating controls will be hidden.
    */
   hideControls?: boolean;
+  /** The session's connection, which the page owns with the document model. */
+  connection: ConnectionStatus;
+  reconnect: () => Promise<void>;
 }
 
-export const EditApp: React.FC<AppProps> = ({
+/** The notebook editor. Rendered inside the page's `DocumentProvider`. */
+export const EditApp: React.FC<EditAppProps> = ({
   userConfig,
   appConfig,
   hideControls = false,
+  connection,
+  reconnect,
 }) => {
-  const { setCells, mergeAllColumns, collapseAllCells, expandAllCells } =
+  const { mergeAllColumns, collapseAllCells, expandAllCells } =
     useCellActions();
   const viewState = useAtomValue(viewStateAtom);
   const numColumns = useAtomValue(numColumnsAtom);
   const filename = useFilename();
-  const setLastSavedNotebook = useSetAtom(lastSavedNotebookAtom);
   const { sendComponentValues, sendInterrupt } = useRequestClient();
 
   const isEditing = viewState.mode === "edit";
@@ -84,18 +90,6 @@ export const EditApp: React.FC<AppProps> = ({
     };
   }, [sendComponentValues]);
 
-  const { connection, reconnect, documentModel } = useNotebookSession({
-    autoInstantiate: userConfig.runtime.auto_instantiate,
-    setCells: (cells, layout) => {
-      setCells(cells);
-      const names = cells.map((cell) => cell.name);
-      const codes = cells.map((cell) => cell.code);
-      const configs = cells.map((cell) => cell.config);
-      setLastSavedNotebook({ names, codes, configs, layout });
-    },
-    sessionId: getSessionId(),
-  });
-
   // Update document title whenever filename or app_title changes
   useEffect(() => {
     // Set document title: app_title takes precedence, then filename, then default
@@ -175,7 +169,7 @@ export const EditApp: React.FC<AppProps> = ({
   );
 
   return (
-    <DocumentProvider model={documentModel}>
+    <>
       <SandboxController onReconnect={reconnect} />
       <AppContainer
         connection={connection}
@@ -229,6 +223,7 @@ export const EditApp: React.FC<AppProps> = ({
           />
         </TooltipProvider>
       )}
-    </DocumentProvider>
+      <LazyCommandPalette />
+    </>
   );
 };

package/src/core/islands/__tests__/bootstrap.test.ts

@@ -1,7 +1,14 @@
 /* Copyright 2026 Marimo. All rights reserved. */
 
+import type { NotebookTransport } from "@marimo-team/client";
 import { createStore } from "jotai";
 import { afterEach, describe, expect, it, vi } from "vitest";
+import { cellId } from "@/__tests__/branded";
+import { notebookAtom } from "@/core/cells/cells";
+import {
+  fakeBackend,
+  fixtureMessages,
+} from "@/core/notebook/replica/__tests__/fixtures";
 import { RuntimeState } from "@/core/kernel/RuntimeState";
 import { initializeIslands } from "../bootstrap";
 import type { IslandsPyodideBridge } from "../bridge";
@@ -32,6 +39,7 @@ describe("initializeIslands", () => {
         consumeMessage = consumer;
       },
       sendComponentValues: vi.fn(),
+      notebookBackend: fakeBackend(),
     } as unknown as IslandsPyodideBridge;
     const store = createStore();
     vi.spyOn(RuntimeState.INSTANCE, "start").mockImplementation(
@@ -59,3 +67,51 @@ describe("initializeIslands", () => {
     expect(store.get(islandsHydratedAtom)).toBe(true);
   });
 });
+
+describe("initializeIslands document", () => {
+  afterEach(vi.restoreAllMocks);
+
+  it("fills the notebook's cells from the managed session's host document", async () => {
+    const root = document.createElement("div");
+    root.innerHTML =
+      '<marimo-island data-app-id="app-1" data-reactive="true"></marimo-island>';
+    let handlers: Parameters<NotebookTransport>[1] | undefined;
+    const transport: NotebookTransport = (_attempt, opened) =>
+      new Promise(() => {
+        handlers = opened;
+        opened.open();
+      });
+    const bridge = {
+      consumeNotifications: vi.fn(),
+      sendComponentValues: vi.fn(),
+      notebookBackend: fakeBackend({ transport }),
+    } as unknown as IslandsPyodideBridge;
+    const store = createStore();
+    vi.spyOn(RuntimeState.INSTANCE, "start").mockImplementation(
+      () => undefined,
+    );
+
+    await initializeIslands({
+      bridge,
+      store,
+      root,
+      autoInitializePlugins: false,
+    });
+    await vi.waitFor(() => expect(handlers).toBeDefined());
+    for (const message of fixtureMessages("notebook-events.sse")) {
+      handlers?.message(message);
+    }
+
+    const notebook = store.get(notebookAtom);
+    expect(notebook.cellIds.inOrderIds).toEqual(["c1", "c3", "k7f2"]);
+    expect(notebook.cellData[cellId("c3")].code).toBe("x = 21");
+    expect(notebook.cellRuntime[cellId("c3")]).toBeDefined();
+
+    handlers?.message({ event: "reset", data: {} });
+    expect(store.get(notebookAtom).cellIds.inOrderIds).toEqual([
+      "c1",
+      "c3",
+      "k7f2",
+    ]);
+  });
+});

package/src/core/islands/__tests__/bridge.test.ts

@@ -396,6 +396,48 @@ describe("IslandsPyodideBridge", () => {
       expect(consumeMessage.mock.calls).toEqual([notification]);
     });
 
+    it("forwards the managed session's host events to the notebook stream", async () => {
+      await initializeSingleApp();
+      const handlers = { open: vi.fn(), message: vi.fn() };
+      void bridge.notebookBackend.transport(
+        { signal: new AbortController().signal, cursor: null },
+        handlers,
+      );
+      expect(handlers.open).toHaveBeenCalledOnce();
+
+      const cell = { event: "cell", id: "1", data: { cell: { id: "c1" } } };
+      sendNotebookEvent(1, cell);
+      sendNotebookEvent(2, { event: "ready", id: "2", data: {} });
+      sendKernelMessage(1);
+
+      expect(handlers.message.mock.calls).toEqual([[cell]]);
+      expect(() => bridge.notebookBackend.notebook("n1")).toThrow(
+        "Islands do not edit",
+      );
+    });
+
+    it("resets the notebook stream when the managed session is replaced or stopped", async () => {
+      await initializeSingleApp();
+      const handlers = { open: vi.fn(), message: vi.fn() };
+      void bridge.notebookBackend.transport(
+        { signal: new AbortController().signal, cursor: null },
+        handlers,
+      );
+
+      mockSingleApp("generated app 2");
+      await bridge.initializeApps();
+      const ready = { event: "ready", id: "2", data: {} };
+      sendNotebookEvent(1, { event: "ready", id: "1", data: {} });
+      sendNotebookEvent(2, ready);
+      await bridge.stopSession("app-1");
+
+      expect(handlers.message.mock.calls).toEqual([
+        [{ event: "reset", data: {} }],
+        [ready],
+        [{ event: "reset", data: {} }],
+      ]);
+    });
+
     it("holds controls until the first app session is ready", async () => {
       setSingleApp();
 

package/src/core/islands/bootstrap.ts

@@ -17,10 +17,14 @@ import {
   notebookAtom,
   notebookReducer,
 } from "../cells/cells";
+import { followDocument } from "../cells/document";
+import { notebookCells } from "../cells/utils";
+import { createDocumentModel } from "../notebook/document";
 import { defineCustomElement } from "../dom/defineCustomElement";
 import { MarimoValueInputEvent } from "../dom/events";
 import { UI_ELEMENT_REGISTRY } from "../dom/uiregistry";
 import { FUNCTIONS_REGISTRY } from "../functions/FunctionRegistry";
+import { createNotebookReplica } from "../notebook/replica";
 import {
   handleCellNotificationeration,
   handleKernelReady,
@@ -124,6 +128,23 @@ export async function initializeIslands(
     handleNotification(notification, actions, store, sessionGeneration);
   });
 
+  // The notebook's cells follow the managed session's host document.
+  const replica = createNotebookReplica(store);
+  replica.connect(bridge.notebookBackend, { onNotification: Functions.NOOP });
+  const documentModel = createDocumentModel({
+    store,
+    documentAtom: replica.documentAtom,
+    backendAtom: replica.backendAtom,
+  });
+  const lifetime = new AbortController();
+  documentModel.start({ signal: lifetime.signal });
+  followDocument({
+    store,
+    model: documentModel,
+    reconcile: actions.reconcileDocument,
+    signal: lifetime.signal,
+  });
+
   // Track first user interaction before initialization completes
   document.addEventListener(
     MarimoValueInputEvent.TYPE,
@@ -196,7 +217,9 @@ function handleNotification(
       case "kernel-ready":
         handleKernelReady(msg, {
           autoInstantiate: true,
-          setCells: actions.setCells,
+          cells: notebookCells(store.get(notebookAtom)),
+          setLastRun: actions.setLastRun,
+          resetCellRuntimes: actions.resetCellRuntimes,
           setLayoutData: Functions.NOOP,
           setAppConfig: Functions.NOOP,
           setCapabilities: Functions.NOOP,

package/src/core/islands/bridge.ts

@@ -1,6 +1,7 @@
 /* Copyright 2026 Marimo. All rights reserved. */
 /* oxlint-disable typescript/no-explicit-any */
 
+import type { NotebookTransport } from "@marimo-team/client";
 import { getWorkerRPC } from "@/core/wasm/rpc";
 import { Deferred } from "@/utils/Deferred";
 import { throwNotImplemented } from "@/utils/functions";
@@ -8,6 +9,7 @@ import { Logger } from "@/utils/Logger";
 import { jsonParseWithSpecialChar } from "@/utils/json/json-parser";
 import { generateUUID } from "@/utils/uuid";
 import { initialNotebookState, notebookAtom } from "../cells/cells";
+import type { NotebookBackend } from "../notebook/replica";
 import type { ClientCommand, NotificationMessage } from "../kernel/messages";
 import type { EditRequests, RunRequests } from "../network/types";
 import { store as defaultStore } from "../state/jotai";
@@ -70,9 +72,37 @@ export class IslandsPyodideBridge implements RunRequests, EditRequests {
   private readonly activeSessionGenerations = new Set<number>();
   private appTransition = Promise.resolve();
   private workerReady = new Deferred<void>();
+  /** Where the managed session's host events go while a replica follows them. */
+  private stream: Parameters<NotebookTransport>[1] | undefined;
 
   public initialized = new Deferred<void>();
 
+  /**
+   * The managed session's host stream, for a replica to follow. Replacing
+   * or stopping the session resets the stream, and the next session's
+   * snapshot follows.
+   */
+  public readonly notebookBackend: NotebookBackend = {
+    transport: ({ signal }, handlers) =>
+      new Promise((resolve) => {
+        this.stream = handlers;
+        handlers.open();
+        signal.addEventListener(
+          "abort",
+          () => {
+            if (this.stream === handlers) {
+              this.stream = undefined;
+            }
+            resolve();
+          },
+          { once: true },
+        );
+      }),
+    notebook: () => {
+      throw new Error("Islands do not edit notebooks");
+    },
+  };
+
   public constructor(config: IslandsBridgeConfig = {}) {
     this.store = config.store || defaultStore;
     this.root = config.root || document;
@@ -119,18 +149,29 @@ export class IslandsPyodideBridge implements RunRequests, EditRequests {
         if (!this.activeSessionGenerations.has(sessionGeneration)) {
           return;
         }
-        // Islands follow only the notebook's marimo notifications.
         const frame = jsonParseWithSpecialChar(event) as {
           event: string;
           data: { marimo?: NotificationMessage };
         };
-        if (frame.event === "marimo" && frame.data.marimo) {
-          this.notificationConsumer?.(frame.data.marimo, sessionGeneration);
+        if (frame.event === "marimo") {
+          if (frame.data.marimo) {
+            this.notificationConsumer?.(frame.data.marimo, sessionGeneration);
+          }
+          return;
+        }
+        // The replica follows the managed session's host document.
+        if (sessionGeneration === this.session?.sessionGeneration) {
+          this.stream?.message(frame);
         }
       },
     );
   }
 
+  /** The managed session is gone, so its stream starts over. */
+  private resetNotebookStream(): void {
+    this.stream?.message({ event: "reset", data: {} });
+  }
+
   public async initializeApps(): Promise<void> {
     await this.enqueueAppTransition(async () => {
       await this.workerReady.promise;
@@ -185,6 +226,7 @@ export class IslandsPyodideBridge implements RunRequests, EditRequests {
       const replacesSession = managesSingleApp && previousSession !== undefined;
       this.activeSessionGenerations.add(request.sessionGeneration);
       if (replacesSession) {
+        this.resetNotebookStream();
         this.store.set(notebookAtom, initialNotebookState());
       }
       if (managesSingleApp || !previousSession) {
@@ -236,6 +278,7 @@ export class IslandsPyodideBridge implements RunRequests, EditRequests {
       this.store.set(islandsPendingInitialRunsAtom, null);
       this.session = undefined;
       this.sessionReady = new Deferred<AppSession>();
+      this.resetNotebookStream();
       this.store.set(notebookAtom, initialNotebookState());
     });
   }
@@ -340,12 +383,10 @@ export class IslandsPyodideBridge implements RunRequests, EditRequests {
   public sendInterrupt = throwNotImplemented;
   public sendShutdown = throwNotImplemented;
   public sendFormat = throwNotImplemented;
-  public sendDeleteCell = throwNotImplemented;
   public sendInstallMissingPackages = throwNotImplemented;
   public sendCodeCompletionRequest = throwNotImplemented;
   public saveUserConfig = throwNotImplemented;
   public saveAppConfig = throwNotImplemented;
-  public saveCellConfig = throwNotImplemented;
   public sendRestart = throwNotImplemented;
   public readCode = throwNotImplemented;
   public readSnippets = throwNotImplemented;