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

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

package/src/core/kernel/__tests__/handlers.test.ts

@@ -1,7 +1,8 @@
 /* Copyright 2026 Marimo. All rights reserved. */
 import { describe, expect, it, vi } from "vitest";
 import { cellId } from "@/__tests__/branded";
-import { buildCellData, buildLayoutState } from "../handlers";
+import { createCell } from "@/core/cells/types";
+import { buildLayoutState } from "../handlers";
 import type { NotificationMessageData } from "../messages";
 import { queryParamHandlers } from "../queryParamHandlers";
 
@@ -63,152 +64,9 @@ describe("queryParamHandlers", () => {
   });
 });
 
-describe("buildCellData", () => {
-  it("should build cell data from kernel-ready data", () => {
-    const kernelReadyData: NotificationMessageData<"kernel-ready"> = {
-      cell_ids: [cellId("cell1"), cellId("cell2")],
-      codes: ["x = 1", "y = 2"],
-      names: ["__", "__"],
-      configs: [
-        { disabled: false, hide_code: false },
-        { disabled: false, hide_code: false },
-      ],
-      layout: null,
-      resumed: false,
-      ui_values: {},
-      last_executed_code: {},
-      last_execution_time: {},
-      app_config: {
-        width: "normal",
-        app_title: null,
-        layout_file: null,
-        css_file: null,
-        auto_download: [],
-      },
-      kiosk: false,
-      capabilities: {
-        terminal: false,
-      },
-      auto_instantiated: false,
-      consumer_capabilities: { edit: true, interact: true },
-    };
-
-    const cells = buildCellData(kernelReadyData);
-
-    expect(cells).toMatchInlineSnapshot(`
-      [
-        {
-          "code": "x = 1",
-          "config": {
-            "disabled": false,
-            "hide_code": false,
-          },
-          "edited": false,
-          "id": "cell1",
-          "lastCodeRun": null,
-          "lastExecutionTime": null,
-          "name": "__",
-          "serializedEditorState": null,
-        },
-        {
-          "code": "y = 2",
-          "config": {
-            "disabled": false,
-            "hide_code": false,
-          },
-          "edited": false,
-          "id": "cell2",
-          "lastCodeRun": null,
-          "lastExecutionTime": null,
-          "name": "__",
-          "serializedEditorState": null,
-        },
-      ]
-    `);
-  });
-
-  it("should mark cells as edited when code differs from last executed code", () => {
-    const kernelReadyData: NotificationMessageData<"kernel-ready"> = {
-      cell_ids: [cellId("cell1"), cellId("cell2")],
-      codes: ["x = 1", "y = 3"],
-      names: ["__", "__"],
-      configs: [
-        { disabled: false, hide_code: false },
-        { disabled: false, hide_code: false },
-      ],
-      layout: null,
-      resumed: false,
-      ui_values: {},
-      last_executed_code: {
-        cell1: "x = 1",
-        cell2: "y = 2",
-      },
-      last_execution_time: {
-        cell1: 1_234_567_890,
-        cell2: 1_234_567_891,
-      },
-      app_config: {
-        width: "normal",
-        app_title: null,
-        layout_file: null,
-        css_file: null,
-        auto_download: [],
-      },
-      kiosk: false,
-      capabilities: {
-        terminal: false,
-      },
-      auto_instantiated: false,
-      consumer_capabilities: { edit: true, interact: true },
-    };
-
-    const cells = buildCellData(kernelReadyData);
-
-    expect(cells[0].edited).toBe(false);
-    expect(cells[1].edited).toBe(true);
-    expect(cells[0].lastCodeRun).toBe("x = 1");
-    expect(cells[1].lastCodeRun).toBe("y = 2");
-  });
-
-  it("should handle empty cell data", () => {
-    const kernelReadyData: NotificationMessageData<"kernel-ready"> = {
-      cell_ids: [],
-      codes: [],
-      names: [],
-      configs: [],
-      layout: null,
-      resumed: false,
-      ui_values: {},
-      last_executed_code: {},
-      last_execution_time: {},
-      app_config: {
-        width: "normal",
-        app_title: null,
-        layout_file: null,
-        css_file: null,
-        auto_download: [],
-      },
-      kiosk: false,
-      capabilities: {
-        terminal: false,
-      },
-      auto_instantiated: false,
-      consumer_capabilities: { edit: true, interact: true },
-    };
-
-    const cells = buildCellData(kernelReadyData);
-
-    expect(cells).toMatchInlineSnapshot("[]");
-  });
-});
-
 describe("buildLayoutState", () => {
   it("should build default layout state when no layout is provided", () => {
     const kernelReadyData: NotificationMessageData<"kernel-ready"> = {
-      cell_ids: [cellId("cell1")],
-      codes: ["x = 1"],
-      names: ["__"],
-      configs: [{ disabled: false, hide_code: false }],
       layout: null,
       resumed: false,
       ui_values: {},
@@ -229,7 +87,7 @@ describe("buildLayoutState", () => {
       consumer_capabilities: { edit: true, interact: true },
     };
 
-    const cells = buildCellData(kernelReadyData);
+    const cells = [createCell({ id: cellId("cell1") })];
     const mockSetLayoutData = vi.fn();
     const layoutState = buildLayoutState(
       kernelReadyData,
@@ -248,13 +106,6 @@ describe("buildLayoutState", () => {
 
   it("should build layout state with vertical layout", () => {
     const kernelReadyData: NotificationMessageData<"kernel-ready"> = {
-      cell_ids: [cellId("cell1"), cellId("cell2")],
-      codes: ["x = 1", "y = 2"],
-      names: ["__", "__"],
-      configs: [
-        { disabled: false, hide_code: false },
-        { disabled: false, hide_code: false },
-      ],
       layout: {
         type: "vertical",
         data: [],
@@ -278,7 +129,10 @@ describe("buildLayoutState", () => {
       consumer_capabilities: { edit: true, interact: true },
     };
 
-    const cells = buildCellData(kernelReadyData);
+    const cells = [
+      createCell({ id: cellId("cell1") }),
+      createCell({ id: cellId("cell2") }),
+    ];
     const mockSetLayoutData = vi.fn();
     const layoutState = buildLayoutState(
       kernelReadyData,

package/src/core/kernel/handlers.ts

@@ -3,7 +3,8 @@ import type { LayoutType } from "@/components/editor/renderers/types";
 import { Logger } from "@/utils/Logger";
 import { Objects } from "@/utils/objects";
 import type { UIElementId } from "../cells/ids";
-import { type CellData, createCell } from "../cells/types";
+import type { CellActions } from "../cells/cells";
+import type { CellData } from "../cells/types";
 import { type AppConfig, AppConfigSchema } from "../config/config-schema";
 import { UI_ELEMENT_REGISTRY } from "../dom/uiregistry";
 import { deserializeLayoutState } from "../layout/layout";
@@ -19,41 +20,6 @@ import type { KernelState } from "./state";
 
 type KernelReadyData = NotificationMessageData<"kernel-ready">;
 
-/**
- * Build cells from kernel-ready data.
- */
-export function buildCellData(data: KernelReadyData): CellData[] {
-  const {
-    codes,
-    names,
-    configs,
-    cell_ids,
-    last_executed_code,
-    last_execution_time,
-  } = data;
-
-  const lastExecutedCode = last_executed_code || {};
-  const lastExecutionTime = last_execution_time || {};
-
-  return codes.map((code, i) => {
-    const cellId = cell_ids[i];
-    const lastCodeRun = lastExecutedCode[cellId];
-
-    // A cell is edited if the code has changed since the last time it was run
-    const edited = lastCodeRun ? lastCodeRun !== code : false;
-
-    return createCell({
-      id: cellId,
-      code,
-      edited,
-      name: names[i],
-      lastCodeRun: lastExecutedCode[cellId] ?? null,
-      lastExecutionTime: lastExecutionTime[cellId] ?? null,
-      config: configs[i],
-    });
-  });
-}
-
 /**
  * Build layout state from kernel-ready data.
  */
@@ -95,7 +61,10 @@ export function handleKernelReady(
   data: KernelReadyData,
   opts: {
     autoInstantiate: boolean;
-    setCells: (cells: CellData[], layout: LayoutState) => void;
+    /** The notebook's cells, as the host document provided them. */
+    cells: CellData[];
+    setLastRun: CellActions["setLastRun"];
+    resetCellRuntimes: CellActions["resetCellRuntimes"];
     setLayoutData: (payload: {
       layoutView: LayoutType;
       data: LayoutData;
@@ -104,17 +73,13 @@ export function handleKernelReady(
     setKernelState: (state: KernelState) => void;
     setAppConfig: (config: AppConfig) => void;
     onError: (error: Error) => void;
-    /**
-     * If provided, these cells will be used instead of the cells from
-     * kernel-ready. This allows preserving local edits made before connecting.
-     */
-    existingCells?: CellData[];
   },
-) {
+): LayoutState {
   const {
-    existingCells,
     autoInstantiate,
-    setCells,
+    cells,
+    setLastRun,
+    resetCellRuntimes,
     setLayoutData,
     setAppConfig,
     setCapabilities,
@@ -124,15 +89,12 @@ export function handleKernelReady(
   const { resumed, ui_values, app_config, capabilities, auto_instantiated } =
     data;
 
-  // Use existing cells if provided (local pre-connect edits), otherwise build from kernel-ready
-  // If the kernel was resumed, we don't want to use the existing cells because they may be stale.
-  const hasExistingCells = existingCells && existingCells.length > 0;
-  const cells =
-    hasExistingCells && !resumed ? existingCells : buildCellData(data);
-
-  // Set up layout and cells
+  resetCellRuntimes();
+  setLastRun({
+    lastCodeRun: data.last_executed_code ?? {},
+    lastExecutionTime: data.last_execution_time ?? {},
+  });
   const layoutState = buildLayoutState(data, cells, setLayoutData);
-  setCells(cells, layoutState);
 
   // Set app config and capabilities
   const parsedAppConfig = AppConfigSchema.safeParse(app_config);
@@ -145,7 +107,7 @@ export function handleKernelReady(
 
   // If the kernel was already instantiated server-side (e.g. run mode), we're done
   if (auto_instantiated) {
-    return;
+    return layoutState;
   }
 
   // If resumed, restore UI element values from kernel and we're done
@@ -156,16 +118,13 @@ export function handleKernelReady(
       const uiObjectId: UIElementId = objectId;
       UI_ELEMENT_REGISTRY.set(uiObjectId, value);
     }
-    return;
+    return layoutState;
   }
 
   // If we already have values for some objects, we should
   // send them to the kernel. This may happen after re-connecting
   // to the kernel after the computer wakes from sleep.
   const { objectIds, values } = collectUIElementValues();
-  const codesToSend = hasExistingCells
-    ? Object.fromEntries(existingCells.map((c) => [c.id, c.code]))
-    : undefined;
 
   // Send instantiate request to kernel
   void getRequestClient()
@@ -173,7 +132,6 @@ export function handleKernelReady(
       objectIds,
       values,
       autoRun: autoInstantiate,
-      codes: codesToSend,
     })
     .then(() => {
       setKernelState({ isInstantiated: true, error: null });
@@ -182,6 +140,7 @@ export function handleKernelReady(
       setKernelState({ isInstantiated: false, error: error });
       onError(new Error("Failed to instantiate", { cause: error }));
     });
+  return layoutState;
 }
 
 export function handleRemoveUIElements(

package/src/core/run-app.tsx

@@ -11,11 +11,7 @@ import { Spinner } from "@/components/icons/spinner";
 import { buttonVariants } from "@/components/ui/button";
 import { cn } from "@/utils/cn";
 import { CellsRenderer } from "../components/editor/renderers/cells-renderer";
-import {
-  hasCellsAtom,
-  notebookIsRunningAtom,
-  useCellActions,
-} from "./cells/cells";
+import { hasCellsAtom, notebookIsRunningAtom } from "./cells/cells";
 import type { AppConfig } from "./config/config-schema";
 import { DocumentProvider } from "./notebook/document";
 import { RuntimeState } from "./kernel/RuntimeState";
@@ -41,7 +37,6 @@ const canPaintRunAppAtom = atom(
 );
 
 export const RunApp: React.FC<AppProps> = ({ appConfig, hideHeader }) => {
-  const { setCells } = useCellActions();
   const { sendComponentValues } = useRequestClient();
 
   // Initialize RuntimeState event-listeners
@@ -54,7 +49,6 @@ export const RunApp: React.FC<AppProps> = ({ appConfig, hideHeader }) => {
 
   const { connection, reconnect, documentModel } = useNotebookSession({
     autoInstantiate: true,
-    setCells: setCells,
     sessionId: getSessionId(),
   });
 

package/src/core/session/__tests__/useNotebookSession.hook.test.tsx

@@ -159,7 +159,6 @@ function renderConnectionHook(
       useNotebookSession({
         sessionId: SESSION_ID,
         autoInstantiate: false,
-        setCells: vi.fn(),
       }),
     { wrapper, reactStrictMode },
   );
@@ -334,10 +333,6 @@ it.each(["kernel-ready", "reconnected"])(
     act(() =>
       conn.send({
         op,
-        cell_ids: [],
-        codes: [],
-        names: [],
-        configs: [],
         layout: null,
         resumed: true,
         ui_values: {},
@@ -367,7 +362,6 @@ describe("connection notice", () => {
       const { reconnect } = useNotebookSession({
         sessionId: SESSION_ID,
         autoInstantiate: false,
-        setCells: vi.fn(),
       });
       return (
         <ConnectionNotice

package/src/core/session/useNotebookSession.tsx

@@ -5,7 +5,8 @@ import { useAtom, useAtomValue, useSetAtom, useStore } from "jotai";
 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 { getNotebook, notebookAtom, useCellActions } from "@/core/cells/cells";
+import { followDocument } from "@/core/cells/document";
 import { AUTOCOMPLETER } from "@/core/codemirror/completion/Autocompleter";
 import { createDocumentModel } from "@/core/notebook/document";
 import type { NotificationPayload } from "@/core/kernel/messages";
@@ -22,14 +23,13 @@ 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 { useDataSourceActions } from "../datasets/data-source-connections";
@@ -75,30 +75,23 @@ import type { VariableName } from "../variables/types";
 import { activeDocumentModelAtom } from "./active-notebook";
 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;
+  /** Called with the layout the kernel reported, after its cells are ready. */
+  onKernelReady?: (layout: LayoutState) => void;
 }) {
-  const { autoInstantiate, sessionId, setCells } = opts;
+  const { autoInstantiate, sessionId, onKernelReady } = opts;
   const { showBoundary } = useErrorBoundary();
+  const store = useStore();
 
-  const { handleCellMessage } = useCellActions();
-  const { reconcileDocument } = useCellActions({ skipMiddleware: true });
+  const {
+    handleCellMessage,
+    reconcileDocument,
+    setLastRun,
+    resetCellRuntimes,
+  } = useCellActions();
   const { addCellNotification } = useRunsActions();
   const setKernelState = useSetAtom(kernelStateAtom);
   const setInitialRunCompleted = useSetAtom(initialRunCompletedAtom);
@@ -153,18 +146,18 @@ 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,
         });
+        onKernelReady?.(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
@@ -393,7 +386,6 @@ export function useNotebookSession(opts: {
     }
   };
 
-  const store = useStore();
   const onMessageError = (error: unknown) => {
     toast({
       title: "Failed to handle message",
@@ -443,16 +435,13 @@ export function useNotebookSession(opts: {
     const controller = new AbortController();
     documentModel.start({ signal: controller.signal });
     store.set(activeDocumentModelAtom, documentModel);
-    const follow = () => {
-      const document = store.get(documentModel.documentAtom);
-      if (document !== null) {
-        reconcileDocument(document);
-      }
-    };
-    const unsubscribe = store.sub(documentModel.documentAtom, follow);
-    follow();
+    followDocument({
+      store,
+      model: documentModel,
+      reconcile: reconcileDocument,
+      signal: controller.signal,
+    });
     return () => {
-      unsubscribe();
       controller.abort();
       if (store.get(activeDocumentModelAtom) === documentModel) {
         store.set(activeDocumentModelAtom, null);