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

package/src/core/session/active-notebook.ts

@@ -0,0 +1,28 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+import { atom } from "jotai";
+import type { DocumentModel } from "../notebook/document";
+import type { NotebookReplica } from "../notebook/replica";
+
+/** The replica owned by the mounted notebook session. */
+export const activeNotebookReplicaAtom = atom<NotebookReplica | null>(null);
+
+/** The document model owned by the mounted notebook session. */
+export const activeDocumentModelAtom = atom<DocumentModel | null>(null);
+
+/** The backend of the connected notebook, or null while none is. */
+export const activeNotebookBackendAtom = atom((get) => {
+  const replica = get(activeNotebookReplicaAtom);
+  return replica === null ? null : get(replica.backendAtom);
+});
+
+/**
+ * The document the host last confirmed, or null while no notebook is
+ * connected. Facts the host owns read this one, such as permission, the
+ * app config, saved state and whether source is available. What the user
+ * is editing reads the model's document, which adds this browser's drafts.
+ */
+export const confirmedDocumentAtom = atom((get) => {
+  const replica = get(activeNotebookReplicaAtom);
+  return replica === null ? null : get(replica.documentAtom);
+});

package/src/core/session/backend.ts

@@ -0,0 +1,123 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+/** The backend the editor follows, chosen once for its mode. */
+
+import { NotebookClient, StreamStopped } from "@marimo-team/client";
+import type { NotebookId } from "@marimo-team/client/protocol";
+import type { NotebookBackend } from "../notebook/replica";
+import { PyodideBridge } from "../wasm/bridge";
+import { isWasm } from "../wasm/utils";
+import { sessionSse, websocket, withHeaders } from "../transports/server";
+import {
+  createStaticNotebookBackend,
+  emptySnapshot,
+  type StaticNotebookSnapshot,
+  startWithSnapshot,
+} from "../transports/static";
+
+/** How the editor follows its notebook. */
+export interface EditorBackendOptions {
+  static: boolean;
+  /** The notebook the page was served with, shown until the host answers. */
+  snapshot: StaticNotebookSnapshot | null;
+  transportType: "websocket" | "sse";
+  /** The session route, read before each connection. */
+  url: () => string;
+  /** The runtime's headers, read before each request. */
+  headers: () => Record<string, string>;
+  /** The host URL of a notebook the stream reported. */
+  notebookUrl: (notebookId: NotebookId) => URL;
+  /** Wait until the host can accept the first connection. */
+  waitUntilReady: () => Promise<void>;
+  /** Whether the host is reachable before an explicit reconnect. */
+  checkHealth: () => Promise<boolean>;
+}
+
+/**
+ * Read the session route once the host has started.
+ *
+ * Every connection waits for the same startup. Once it settles, even by
+ * giving up, connections read the route directly. An explicit reconnect
+ * checks the host's health instead.
+ */
+function afterStartup({
+  url,
+  waitUntilReady,
+}: Pick<EditorBackendOptions, "url" | "waitUntilReady">) {
+  let startup: Promise<void> | undefined;
+  let settled = false;
+  return async () => {
+    if (!settled) {
+      startup ??= waitUntilReady().finally(() => {
+        settled = true;
+      });
+      try {
+        await startup;
+      } catch (error) {
+        throw new StreamStopped("Healthy connection never made", {
+          cause: error,
+        });
+      }
+    }
+    return url();
+  };
+}
+
+/**
+ * Follow the notebook running in the WASM worker.
+ *
+ * The worker's backend exists once the worker has started, so the first
+ * connection creates it. Commands are only for a notebook its stream reported.
+ */
+function wasm(): NotebookBackend {
+  let backend: Promise<NotebookBackend> | undefined;
+  let started: NotebookBackend | undefined;
+  return {
+    transport: async (attempt, handlers) => {
+      backend ??= PyodideBridge.INSTANCE.createNotebookBackend();
+      try {
+        started = await backend;
+      } catch (error) {
+        throw new StreamStopped("The WASM notebook did not start", {
+          cause: error,
+        });
+      }
+      return started.transport(attempt, handlers);
+    },
+    notebook: (id) => {
+      if (started === undefined) {
+        throw new Error("The WASM notebook has not started");
+      }
+      return started.notebook(id);
+    },
+  };
+}
+
+/** Choose the backend for the editor's notebook. */
+export function createEditorBackend(
+  options: EditorBackendOptions,
+): NotebookBackend {
+  if (options.static) {
+    return createStaticNotebookBackend(
+      options.snapshot ?? emptySnapshot("static"),
+    );
+  }
+  const live = isWasm() ? wasm() : server(options);
+  return options.snapshot ? startWithSnapshot(options.snapshot, live) : live;
+}
+
+/** Follow the notebook on the server, over SSE or a WebSocket. */
+function server(options: EditorBackendOptions): NotebookBackend {
+  const url = afterStartup(options);
+  return {
+    transport:
+      options.transportType === "sse"
+        ? sessionSse(url, options.headers)
+        : websocket(url),
+    notebook: (id) =>
+      new NotebookClient(options.notebookUrl(id), {
+        fetch: withHeaders(options.headers),
+      }),
+    checkHealth: options.checkHealth,
+  };
+}

package/src/core/session/embedded.ts

@@ -0,0 +1,7 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+import { atom } from "jotai";
+import type { StaticNotebookSnapshot } from "../transports/static";
+
+/** The notebook the page was served with, read at mount. Null when none was. */
+export const embeddedSnapshotAtom = atom<StaticNotebookSnapshot | null>(null);

package/src/core/{websocket → session}/types.ts

@@ -1,12 +1,12 @@
 /* Copyright 2026 Marimo. All rights reserved. */
 
-import type { ConnectionStatus as HostConnectionStatus } from "../host/connection";
+import type { ConnectionStatus as HostConnectionStatus } from "../notebook/replica";
 import type { NotificationMessageData } from "../kernel/messages";
 
 export {
   ConnectionState,
   ConnectionClosedReason as WebSocketClosedReason,
-} from "../host/connection";
+} from "../notebook/replica";
 
 export type ConnectionPhase =
   | NotificationMessageData<"startup-progress">["phase"]

package/src/core/session/useNotebookReplica.ts

@@ -0,0 +1,74 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+/** The editor session's replica, following its backend while mounted. */
+
+import { useStore } from "jotai";
+import { useEffect, useMemo } from "react";
+import useEvent from "react-use-event-hook";
+import type {
+  NotificationMessage,
+  NotificationPayload,
+} from "../kernel/messages";
+import {
+  createNotebookReplica,
+  type NotebookReplica,
+} from "../notebook/replica";
+import { activeNotebookReplicaAtom } from "./active-notebook";
+import { createEditorBackend, type EditorBackendOptions } from "./backend";
+
+type Options = EditorBackendOptions & {
+  onNotification: (notification: NotificationPayload) => void;
+};
+
+export function useNotebookReplica(options: Options): NotebookReplica {
+  const store = useStore();
+  const { static: isStatic, snapshot, transportType } = options;
+  const replica = useMemo(() => createNotebookReplica(store), [store]);
+  const url = useEvent(options.url);
+  const headers = useEvent(options.headers);
+  const notebookUrl = useEvent(options.notebookUrl);
+  const waitUntilReady = useEvent(options.waitUntilReady);
+  const checkHealth = useEvent(options.checkHealth);
+  const onNotification = useEvent(options.onNotification);
+  // Opens nothing until the replica follows it.
+  const backend = useMemo(
+    () =>
+      createEditorBackend({
+        static: isStatic,
+        snapshot,
+        transportType,
+        url,
+        headers,
+        notebookUrl,
+        waitUntilReady,
+        checkHealth,
+      }),
+    [
+      isStatic,
+      snapshot,
+      transportType,
+      url,
+      headers,
+      notebookUrl,
+      waitUntilReady,
+      checkHealth,
+    ],
+  );
+
+  useEffect(() => {
+    store.set(activeNotebookReplicaAtom, replica);
+    const disconnect = replica.connect(backend, {
+      // The host validates the envelope, not the native notification contents.
+      onNotification: (notification) =>
+        onNotification({ data: notification as NotificationMessage }),
+    });
+    return () => {
+      disconnect();
+      if (store.get(activeNotebookReplicaAtom) === replica) {
+        store.set(activeNotebookReplicaAtom, null);
+      }
+    };
+  }, [replica, backend, store, onNotification]);
+
+  return replica;
+}

package/src/core/{websocket → session}/useNotebookSession.tsx

@@ -2,17 +2,19 @@
 
 import { invalidatePackageData } from "@/core/packages/package-data";
 import { useAtom, useAtomValue, useSetAtom, useStore } from "jotai";
-import { useLayoutEffect } from "react";
+import { useEffect, useLayoutEffect, useMemo } from "react";
 import { useErrorBoundary } from "react-error-boundary";
 import { toast } from "@/components/ui/use-toast";
-import { getNotebook, useCellActions } from "@/core/cells/cells";
-import { applyTransactionChanges } from "@/core/cells/document-changes";
+import {
+  getNotebook,
+  notebookAtom,
+  useReducerActions,
+} from "@/core/cells/cells";
+import { followDocument } from "@/core/cells/document";
 import { AUTOCOMPLETER } from "@/core/codemirror/completion/Autocompleter";
-import type {
-  NotificationMessageData,
-  NotificationPayload,
-} from "@/core/kernel/messages";
-import { useNotebookClient } from "@/core/websocket/useNotebookClient";
+import { createDocumentModel } from "@/core/notebook/document";
+import type { NotificationPayload } from "@/core/kernel/messages";
+import { useNotebookReplica } from "@/core/session/useNotebookReplica";
 import { renderHTML } from "@/plugins/core/RenderHTML";
 import {
   handleWidgetMessage,
@@ -25,16 +27,15 @@ import { Logger } from "@/utils/Logger";
 import { reloadSafe } from "@/utils/reload-safe";
 import { useAlertActions } from "../alerts/state";
 import { cacheInfoAtom } from "../cache/requests";
-import { SCRATCH_CELL_ID } from "../cells/ids";
 import { useRunsActions } from "../cells/runs";
 import { focusAndScrollCellOutputIntoView } from "../cells/scrollCellIntoView";
 import {
   resyncBreakpoints,
   setActiveLine,
 } from "../codemirror/cells/debugger-state";
-import type { CellData } from "../cells/types";
+import { notebookCells } from "../cells/utils";
 import { capabilitiesAtom } from "../config/capabilities";
-import { connectionTransportTypeAtom, useSetAppConfig } from "../config/config";
+import { connectionTransportTypeAtom } from "../config/config";
 import { useDataSourceActions } from "../datasets/data-source-connections";
 import { invalidateDataSourceDiscovery } from "../datasets/data-source-discovery";
 import type { ConnectionName } from "../datasets/engines";
@@ -61,11 +62,12 @@ import {
   stableSessionIdAtom,
 } from "../kernel/session";
 import { initialRunCompletedAtom, kernelStateAtom } from "../kernel/state";
-import { type LayoutState, useLayoutActions } from "../layout/state";
+import { useLayoutActions } from "../layout/state";
 import { kioskModeAtom } from "../mode";
 import { connectionAtom, startupProgressAtom } from "../network/connection";
+import { savedLayoutAtom } from "../saving/state";
 import type { RequestId } from "../network/DeferredRequestRegistry";
-import { runtimeConfigAtom, useRuntimeManager } from "../runtime/config";
+import { useRuntimeManager } from "../runtime/config";
 import { SECRETS_REGISTRY } from "../secrets/request-registry";
 import { isStaticNotebook } from "../static/static-state";
 import {
@@ -75,47 +77,28 @@ import {
 import { useStorageActions } from "../storage/state";
 import { useVariablesActions } from "../variables/state";
 import type { VariableName } from "../variables/types";
-import { isWasm } from "../wasm/utils";
+import { activeDocumentModelAtom } from "./active-notebook";
+import { embeddedSnapshotAtom } from "./embedded";
 import { WebSocketClosedReason, ConnectionState } from "./types";
 
-const SUPPORTS_LAZY_KERNELS = true;
-
-function getExistingCells(): CellData[] | undefined {
-  if (!SUPPORTS_LAZY_KERNELS) {
-    return undefined;
-  }
-
-  // Remove scratch pad
-  return Object.values(getNotebook().cellData).filter(
-    (cell) => cell.id !== SCRATCH_CELL_ID,
-  );
-}
-
 /** Coordinate session startup and apply notifications to the editor. */
 export function useNotebookSession(opts: {
   sessionId: SessionId;
   autoInstantiate: boolean;
-  setCells: (cells: CellData[], layout: LayoutState) => void;
 }) {
-  const { autoInstantiate, sessionId, setCells } = opts;
+  const { autoInstantiate, sessionId } = opts;
   const { showBoundary } = useErrorBoundary();
+  const store = useStore();
 
-  const { handleCellMessage } = useCellActions();
-  const actionsWithoutMiddleware = useCellActions({ skipMiddleware: true });
-
-  const handleDocumentTransaction = (
-    transaction: NotificationMessageData<"notebook-document-transaction">["transaction"],
-  ) => {
-    applyTransactionChanges(
-      transaction.changes,
-      actionsWithoutMiddleware,
-      () => getNotebook().cellIds.inOrderIds,
-    );
-  };
+  const {
+    handleCellMessage,
+    reconcileDocument,
+    setLastRun,
+    resetCellRuntimes,
+  } = useReducerActions();
   const { addCellNotification } = useRunsActions();
   const setKernelState = useSetAtom(kernelStateAtom);
   const setInitialRunCompleted = useSetAtom(initialRunCompletedAtom);
-  const setAppConfig = useSetAppConfig();
   const { setVariables, setMetadata } = useVariablesActions();
   const { addColumnPreview } = useDatasetsActions();
   const { addDatasets, filterDatasetsFromVariables } = useDatasetsActions();
@@ -135,7 +118,6 @@ export function useNotebookSession(opts: {
   const setStableSessionId = useSetAtom(stableSessionIdAtom);
   const setCapabilities = useSetAtom(capabilitiesAtom);
   const runtimeManager = useRuntimeManager();
-  const { notebookId } = useAtomValue(runtimeConfigAtom);
   const transportType = useAtomValue(connectionTransportTypeAtom);
   const setCacheInfo = useSetAtom(cacheInfoAtom);
   const setKernelStartupError = useSetAtom(kernelStartupErrorAtom);
@@ -167,18 +149,17 @@ export function useNotebookSession(opts: {
         setInitialRunCompleted(
           Boolean(msg.data.resumed || msg.data.auto_instantiated),
         );
-        const existingCells = getExistingCells();
-
-        handleKernelReady(msg.data, {
+        const layout = handleKernelReady(msg.data, {
           autoInstantiate,
-          setCells,
+          cells: notebookCells(store.get(notebookAtom)),
+          setLastRun,
+          resetCellRuntimes,
           setLayoutData,
-          setAppConfig,
           setCapabilities,
           setKernelState,
           onError: showBoundary,
-          existingCells,
         });
+        store.set(savedLayoutAtom, layout);
         if (!msg.data.resumed) {
           // A freshly started kernel may expose a different environment
           // (e.g. new env vars); re-run discovery instead of reusing stale
@@ -265,13 +246,18 @@ export function useNotebookSession(opts: {
         filterStorageFromVariables(variableNames);
         return;
       case "variable-values":
-        setMetadata(
-          msg.data.variables.map((v) => ({
+        setMetadata([
+          ...(msg.data.removed ?? []).map((v) => ({
+            name: v.name as VariableName,
+            dataType: undefined,
+            value: undefined,
+          })),
+          ...msg.data.variables.map((v) => ({
             name: v.name as VariableName,
             dataType: v.datatype,
             value: v.value,
           })),
-        );
+        ]);
         return;
       case "alert":
         toast({
@@ -392,7 +378,7 @@ export function useNotebookSession(opts: {
         );
         return;
       case "notebook-document-transaction":
-        handleDocumentTransaction(msg.data.transaction);
+        // Edits reach the notebook through the host document.
         return;
       case "consumer-capabilities":
         setKioskMode(!msg.data.consumer_capabilities.edit);
@@ -402,7 +388,6 @@ export function useNotebookSession(opts: {
     }
   };
 
-  const store = useStore();
   const onMessageError = (error: unknown) => {
     toast({
       title: "Failed to handle message",
@@ -410,25 +395,24 @@ export function useNotebookSession(opts: {
       variant: "danger",
     });
   };
-  const notebook = useNotebookClient({
-    sessionId,
+  const replica = useNotebookReplica({
     static: isStaticNotebook(),
+    snapshot: useAtomValue(embeddedSnapshotAtom),
     transportType,
-    messageFormat:
-      notebookId && transportType === "websocket" ? "host" : "session",
     url: () =>
       transportType === "sse"
         ? runtimeManager.getSseURL(sessionId).toString()
         : runtimeManager.getWsURL(sessionId).toString(),
     // SSE uses headers for auth. WebSockets carry the token in the URL.
     headers: () => runtimeManager.headers(),
+    notebookUrl: (notebookId) => runtimeManager.getNotebookURL(notebookId),
+    // Only the host's transports read the URL, so static and WASM never wait.
     waitUntilReady: async () => {
-      if (!isStaticNotebook() && !isWasm() && !runtimeManager.isSameOrigin) {
+      if (!runtimeManager.isSameOrigin) {
         await runtimeManager.waitForHealthy();
       }
     },
     checkHealth: () => runtimeManager.reconcileFromHealth(),
-    onMessageError,
     onNotification: (notification) => {
       try {
         handleNotification(notification);
@@ -439,10 +423,39 @@ export function useNotebookSession(opts: {
     },
   });
 
+  const documentModel = useMemo(
+    () =>
+      createDocumentModel({
+        store,
+        documentAtom: replica.documentAtom,
+        backendAtom: replica.backendAtom,
+      }),
+    [store, replica],
+  );
+  // Edits outlive reconnects and end with the session. The notebook's cells
+  // follow the document for as long as the session is mounted.
+  useEffect(() => {
+    const controller = new AbortController();
+    documentModel.start({ signal: controller.signal });
+    store.set(activeDocumentModelAtom, documentModel);
+    followDocument({
+      store,
+      model: documentModel,
+      reconcile: reconcileDocument,
+      signal: controller.signal,
+    });
+    return () => {
+      controller.abort();
+      if (store.get(activeDocumentModelAtom) === documentModel) {
+        store.set(activeDocumentModelAtom, null);
+      }
+    };
+  }, [documentModel, store, reconcileDocument]);
+
   // Subscribe before passive connection setup: static/WASM transports can emit synchronously.
   useLayoutEffect(() => {
     const updateConnection = () => {
-      const status = store.get(notebook.connectionAtom);
+      const status = store.get(replica.connectionAtom);
       if (status.state === ConnectionState.NOT_STARTED) {
         return;
       }
@@ -473,16 +486,16 @@ export function useNotebookSession(opts: {
         return status;
       });
     };
-    const unsubscribe = store.sub(notebook.connectionAtom, updateConnection);
+    const unsubscribe = store.sub(replica.connectionAtom, updateConnection);
     updateConnection();
     return unsubscribe;
   }, [
-    notebook,
+    replica,
     store,
     setConnection,
     setKernelStartupError,
     updateStartupProgress,
   ]);
 
-  return { notebook, connection, reconnect: notebook.reconnect };
+  return { replica, documentModel, connection, reconnect: replica.reconnect };
 }

package/src/core/state/__tests__/jotai.test.tsx

@@ -0,0 +1,108 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+import { act, renderHook } from "@testing-library/react";
+import { atom, Provider } from "jotai";
+import type React from "react";
+import { describe, expect, it, vi } from "vitest";
+import { store, useJotaiEffect, waitFor } from "../jotai";
+
+const wrapper = ({ children }: { children: React.ReactNode }) => (
+  <Provider store={store}>{children}</Provider>
+);
+
+describe("waitFor function", () => {
+  it("should resolve when the atom satisfies the predicate", async () => {
+    const testAtom = atom(0);
+    store.set(testAtom, 10);
+
+    const result = await waitFor(testAtom, (value) => value === 10);
+    expect(result).toBe(10);
+  });
+
+  it("should resolve when the atom changes to satisfy the predicate", async () => {
+    const testAtom = atom(0);
+
+    setTimeout(() => store.set(testAtom, 15), 5);
+    setTimeout(() => store.set(testAtom, 20), 10);
+
+    const result = await waitFor(testAtom, (value) => value === 20);
+    expect(result).toBe(20);
+  });
+});
+
+describe("useJotaiEffect", () => {
+  it("calls the effect with the new and previous values", () => {
+    const testAtom = atom(0);
+    const effect = vi.fn();
+    renderHook(() => useJotaiEffect(testAtom, effect), { wrapper });
+
+    act(() => store.set(testAtom, 1));
+    act(() => store.set(testAtom, 2));
+
+    expect(effect.mock.calls).toEqual([
+      [1, 0],
+      [2, 1],
+    ]);
+  });
+
+  it("calls only the latest callback after a rerender", () => {
+    const testAtom = atom(0);
+    const first = vi.fn();
+    const second = vi.fn();
+    const { rerender } = renderHook(
+      ({ effect }) => useJotaiEffect(testAtom, effect),
+      { initialProps: { effect: first }, wrapper },
+    );
+
+    rerender({ effect: second });
+    act(() => store.set(testAtom, 1));
+
+    expect(first).not.toHaveBeenCalled();
+    expect(second).toHaveBeenCalledTimes(1);
+  });
+
+  it("subscribes once per rerender with an inline callback", () => {
+    const testAtom = atom(0);
+    const calls: number[] = [];
+    const { rerender } = renderHook(
+      () => useJotaiEffect(testAtom, (value) => calls.push(value)),
+      { wrapper },
+    );
+
+    rerender();
+    rerender();
+    act(() => store.set(testAtom, 1));
+
+    expect(calls).toEqual([1]);
+  });
+
+  it("stops listening to the previous atom when the atom changes", () => {
+    const firstAtom = atom(0);
+    const secondAtom = atom(0);
+    const effect = vi.fn();
+    const { rerender } = renderHook(
+      ({ target }) => useJotaiEffect(target, effect),
+      { initialProps: { target: firstAtom }, wrapper },
+    );
+
+    rerender({ target: secondAtom });
+    act(() => store.set(firstAtom, 1));
+    expect(effect).not.toHaveBeenCalled();
+
+    act(() => store.set(secondAtom, 1));
+    expect(effect).toHaveBeenCalledTimes(1);
+  });
+
+  it("unsubscribes on unmount", () => {
+    const testAtom = atom(0);
+    const effect = vi.fn();
+    const { unmount } = renderHook(() => useJotaiEffect(testAtom, effect), {
+      wrapper,
+    });
+
+    unmount();
+    act(() => store.set(testAtom, 1));
+
+    expect(effect).not.toHaveBeenCalled();
+  });
+});

package/src/core/state/jotai.ts

@@ -42,7 +42,7 @@ export function useJotaiEffect<T>(
   const store = useStore();
   useEffect(() => {
     let prevValue = store.get(atom);
-    store.sub(atom, () => {
+    return store.sub(atom, () => {
       const value = store.get(atom);
       effect(value, prevValue);
       prevValue = value;