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

package/dist/{xychartDiagram-S5SC5T6Z-MxyaWjjq.js → xychartDiagram-S5SC5T6Z-1hIHN6Rl.js}

@@ -9,7 +9,7 @@ import { n as s } from "./chunk-Y2CYZVJY-DKKS1VyV.js";
 import { D as c, H as l, K as u, U as d, a as f, b as p, c as m, f as h, v as g, w as _, y as v, z as y } from "./chunk-DU6HZSFF-BwiyaNpg.js";
 import "./chunk-PWAF6VOD-CQw5qATZ.js";
 import { t as b } from "./chunk-GMAD6QVW-BhaJx0X3.js";
-import { p as x } from "./mermaid-CSvw7c71.js";
+import { p as x } from "./mermaid-CZ6FrClS.js";
 //#region ../node_modules/.pnpm/d3-scale@4.0.2/node_modules/d3-scale/src/band.js
 function S() {
 	var e = t().unknown(void 0), r = e.domain, i = e.range, o = 0, s = 1, c, l, u = !1, d = 0, f = 0, p = .5;

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@marimo-team/islands",
-  "version": "0.25.3-dev170",
+  "version": "0.25.3-dev173",
   "main": "dist/main.js",
   "types": "dist/index.d.ts",
   "type": "module",

package/src/core/cells/__tests__/cells.test.ts

@@ -7,6 +7,7 @@ import { createStore } from "jotai";
 import {
   afterAll,
   beforeAll,
+  afterEach,
   beforeEach,
   describe,
   expect,
@@ -3591,31 +3592,9 @@ describe("createTracebackInfoAtom", () => {
   });
 });
 
-describe("setCells snapshot preservation", () => {
+describe("kernel-ready", () => {
   const CELL_A = cellId("A");
   const CELL_B = cellId("B");
-  const newCells: CellData[] = [
-    {
-      id: CELL_A,
-      name: "a",
-      code: "1",
-      edited: false,
-      lastCodeRun: null,
-      lastExecutionTime: null,
-      config: { hide_code: false, disabled: false, column: null },
-      serializedEditorState: null,
-    },
-    {
-      id: CELL_B,
-      name: "b",
-      code: "2",
-      edited: false,
-      lastCodeRun: null,
-      lastExecutionTime: null,
-      config: { hide_code: false, disabled: false, column: null },
-      serializedEditorState: null,
-    },
-  ];
 
   const hydratedState = () =>
     MockNotebook.notebookState({
@@ -3645,44 +3624,97 @@ describe("setCells snapshot preservation", () => {
       },
     });
 
-  beforeEach(async () => {
-    const { isWasm } = await import("@/core/wasm/utils");
-    vi.mocked(isWasm).mockReturnValue(true);
+  it("setLastRun records what the kernel last ran", () => {
+    const state = hydratedState();
+    const next = exportedForTesting.reducer(state, {
+      type: "setLastRun",
+      payload: {
+        lastCodeRun: { [CELL_A]: "0" },
+        lastExecutionTime: { [CELL_A]: 5 },
+      },
+    });
+
+    expect(next.cellData[CELL_A]).toMatchObject({
+      lastCodeRun: "0",
+      lastExecutionTime: 5,
+      edited: true,
+    });
+    expect(next.cellData[CELL_B]).toMatchObject({
+      lastCodeRun: null,
+      lastExecutionTime: null,
+      edited: false,
+    });
+    expect(next.cellRuntime).toBe(state.cellRuntime);
   });
 
-  it("preserves hydrated output in WASM", () => {
+  it("resetCellRuntimes starts every cell's runtime over", () => {
     const next = exportedForTesting.reducer(hydratedState(), {
-      type: "setCells",
-      payload: newCells,
+      type: "resetCellRuntimes",
+      payload: undefined,
     });
 
-    expect(next.cellRuntime[CELL_A].output).toMatchObject({
-      data: "hydrated-A",
+    expect(next.cellRuntime[CELL_A].output).toBeNull();
+    expect(next.cellRuntime[CELL_B].consoleOutputs).toEqual([]);
+  });
+
+  it("resetCellRuntimes starts the scratch cell over too", () => {
+    const state = MockNotebook.notebookState({
+      cellData: { [CELL_A]: { id: CELL_A, code: "1" } },
+      cellRuntime: {
+        [SCRATCH_CELL_ID]: {
+          output: {
+            channel: "output",
+            mimetype: "text/plain",
+            data: "scratch",
+            timestamp: 0,
+          },
+        },
+      },
+    });
+    const next = exportedForTesting.reducer(state, {
+      type: "resetCellRuntimes",
+      payload: undefined,
     });
+
+    expect(next.cellRuntime[SCRATCH_CELL_ID].output).toBeNull();
   });
 
-  it("preserves console-only hydration in WASM", () => {
-    const next = exportedForTesting.reducer(hydratedState(), {
-      type: "setCells",
-      payload: newCells,
+  describe("resetCellRuntimes in WASM", () => {
+    beforeEach(async () => {
+      const { isWasm } = await import("@/core/wasm/utils");
+      vi.mocked(isWasm).mockReturnValue(true);
     });
 
-    expect(next.cellRuntime[CELL_B].consoleOutputs).toHaveLength(1);
-    expect(next.cellRuntime[CELL_B].consoleOutputs[0]).toMatchObject({
-      data: "hydrated-B-stdout",
+    afterEach(async () => {
+      const { isWasm } = await import("@/core/wasm/utils");
+      vi.mocked(isWasm).mockReturnValue(false);
     });
-  });
 
-  it("resets cells with no prior runtime even in WASM", () => {
-    const empty = MockNotebook.notebookState({ cellData: {} });
-    const next = exportedForTesting.reducer(empty, {
-      type: "setCells",
-      payload: newCells,
+    it("keeps hydrated outputs and console", () => {
+      const next = exportedForTesting.reducer(hydratedState(), {
+        type: "resetCellRuntimes",
+        payload: undefined,
+      });
+
+      expect(next.cellRuntime[CELL_A].output).toMatchObject({
+        data: "hydrated-A",
+      });
+      expect(next.cellRuntime[CELL_B].consoleOutputs[0]).toMatchObject({
+        data: "hydrated-B-stdout",
+      });
     });
 
-    expect(next.cellRuntime[CELL_A].output).toBeNull();
-    expect(next.cellRuntime[CELL_A].consoleOutputs).toEqual([]);
-    expect(next.cellRuntime[CELL_B].output).toBeNull();
-    expect(next.cellRuntime[CELL_B].consoleOutputs).toEqual([]);
+    it("starts cells with no prior runtime over", () => {
+      const state = MockNotebook.notebookState({
+        cellData: { [CELL_A]: { id: CELL_A, code: "1" } },
+      });
+      const next = exportedForTesting.reducer(state, {
+        type: "resetCellRuntimes",
+        payload: undefined,
+      });
+
+      expect(next.cellRuntime[CELL_A].output).toBeNull();
+      expect(next.cellRuntime[CELL_A].consoleOutputs).toEqual([]);
+    });
   });
 });

package/src/core/cells/__tests__/document.test.ts

@@ -29,7 +29,12 @@ import {
   type NotebookState,
   notebookReducer,
 } from "../cells";
-import { applyDocument, flushDocument, withDocumentSave } from "../document";
+import {
+  applyDocument,
+  flushDocument,
+  followDocument,
+  withDocumentSave,
+} from "../document";
 import { type CellId, SCRATCH_CELL_ID } from "../ids";
 
 vi.mock("../../rtc/state", async (importOriginal) => ({
@@ -349,9 +354,14 @@ describe("documentEditMiddleware", () => {
     expect(calls).toEqual([["updateConfig", c, { column: null }]]);
   });
 
-  it("ignores the kernel's cells and runtime state", () => {
+  it("ignores the document and the kernel's run state", () => {
     actions.prepareForRun({ cellId: a });
-    actions.setCells([state.cellData[a]]);
+    actions.setLastRun({
+      lastCodeRun: { [a]: "a = 0" },
+      lastExecutionTime: {},
+    });
+    actions.resetCellRuntimes();
+    actions.reconcileDocument(documentOf([cell(a)]));
 
     expect(calls).toEqual([]);
   });
@@ -367,6 +377,30 @@ describe("documentEditMiddleware", () => {
   });
 });
 
+describe("followDocument", () => {
+  it("reconciles the document now and on each change, until aborted", () => {
+    const documentAtom = atom<EditableDocument | null>(null);
+    const reconcile = vi.fn();
+    const controller = new AbortController();
+
+    followDocument({
+      store,
+      model: { documentAtom },
+      reconcile,
+      signal: controller.signal,
+    });
+    expect(reconcile).not.toHaveBeenCalled();
+
+    const first = documentOf([cell(a)]);
+    store.set(documentAtom, first);
+    expect(reconcile).toHaveBeenCalledWith(first);
+
+    controller.abort();
+    store.set(documentAtom, documentOf([cell(b)]));
+    expect(reconcile).toHaveBeenCalledOnce();
+  });
+});
+
 describe("flushing before host commands", () => {
   const model = {
     syncStatusAtom: atom<DocumentSyncStatus>({ status: "idle" }),

package/src/core/cells/cells.ts

@@ -44,6 +44,7 @@ import {
   createCell,
   createCellConfig,
   createCellRuntimeState,
+  isEdited,
 } from "./types";
 import {
   canUndoDeletes,
@@ -936,12 +937,7 @@ const {
       cellLogs: [...state.cellLogs, ...action.logs],
     };
   },
-  /**
-   * Follow the host document's cells and order.
-   *
-   * Dispatched without middleware, since the document already has these
-   * changes.
-   */
+  /** Follow the host document's cells and order. */
   reconcileDocument: (state, document: EditableDocument) => {
     const nextState = applyDocument(state, document);
     for (const cellId of state.cellIds.inOrderIds) {
@@ -985,36 +981,44 @@ const {
       },
     });
   },
-  setCells: (state, cells: CellData[]) => {
-    const cellData = Object.fromEntries(cells.map((cell) => [cell.id, cell]));
-
-    // WASM has no server-side SessionView to replay outputs, so the
-    // snapshot hydrated by notebookStateFromSession is the only source.
+  /** Record what the kernel last ran for each cell. */
+  setLastRun: (
+    state,
+    action: {
+      lastCodeRun: Record<CellId, string>;
+      lastExecutionTime: Record<CellId, number>;
+    },
+  ) => {
+    const cellData = { ...state.cellData };
+    for (const cellId of state.cellIds.inOrderIds) {
+      const cell = state.cellData[cellId];
+      const lastCodeRun = action.lastCodeRun[cellId] ?? null;
+      cellData[cellId] = {
+        ...cell,
+        lastCodeRun,
+        lastExecutionTime: action.lastExecutionTime[cellId] ?? null,
+        // A cell the kernel never ran is not stale on load.
+        edited: lastCodeRun !== null && isEdited(lastCodeRun, cell.code),
+      };
+    }
+    return { ...state, cellData };
+  },
+  /**
+   * Start each cell's runtime over for a new kernel. WASM keeps the outputs
+   * its mount-time snapshot gave it.
+   */
+  resetCellRuntimes: (state) => {
     const preserveSnapshot = isWasm();
-    const runtimeFor = (cellId: CellId): CellRuntimeState => {
-      if (!preserveSnapshot) {
-        return createCellRuntimeState();
-      }
-      const prev = state.cellRuntime[cellId];
+    const cellRuntime = { ...state.cellRuntime };
+    for (const cellId of [...state.cellIds.inOrderIds, SCRATCH_CELL_ID]) {
+      const previous = state.cellRuntime[cellId];
       const hasSnapshot =
-        prev && (prev.output != null || prev.consoleOutputs.length > 0);
-      return hasSnapshot ? prev : createCellRuntimeState();
-    };
-    const cellRuntime = Object.fromEntries(
-      cells.map((cell) => [cell.id, runtimeFor(cell.id)]),
-    );
-
-    return withScratchCell({
-      ...state,
-      cellIds: MultiColumn.fromIdsAndColumns(
-        cells.map((cell) => [cell.id, cell.config.column]),
-      ),
-      cellData: cellData,
-      cellRuntime: cellRuntime,
-      cellHandles: Object.fromEntries(
-        cells.map((cell) => [cell.id, createRef()]),
-      ),
-    });
+        preserveSnapshot &&
+        previous !== undefined &&
+        (previous.output != null || previous.consoleOutputs.length > 0);
+      cellRuntime[cellId] = hasSnapshot ? previous : createCellRuntimeState();
+    }
+    return { ...state, cellRuntime };
   },
   /**
    * Move focus to next cell

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";
@@ -165,6 +167,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 +320,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/edit-app.tsx

@@ -5,7 +5,7 @@ import { Tooltip } from "radix-ui";
 
 const TooltipProvider = Tooltip.Provider;
 
-import { useAtomValue, useSetAtom } from "jotai";
+import { useAtomValue, useSetAtom, useStore } from "jotai";
 import { useEffect } from "react";
 import { NotStartedConnectionAlert } from "@/components/editor/alerts/connecting-alert";
 import { Controls } from "@/components/editor/controls/Controls";
@@ -29,10 +29,12 @@ import { CellsRenderer } from "../components/editor/renderers/cells-renderer";
 import { useHotkey } from "../hooks/useHotkey";
 import {
   hasCellsAtom,
+  notebookAtom,
   notebookIsRunningAtom,
   numColumnsAtom,
   useCellActions,
 } from "./cells/cells";
+import { notebookCells } from "./cells/utils";
 import type { AppConfig, UserConfig } from "./config/config-schema";
 import { DocumentProvider } from "./notebook/document";
 import { RuntimeState } from "./kernel/RuntimeState";
@@ -64,8 +66,9 @@ export const EditApp: React.FC<AppProps> = ({
   appConfig,
   hideControls = false,
 }) => {
-  const { setCells, mergeAllColumns, collapseAllCells, expandAllCells } =
+  const { mergeAllColumns, collapseAllCells, expandAllCells } =
     useCellActions();
+  const store = useStore();
   const viewState = useAtomValue(viewStateAtom);
   const numColumns = useAtomValue(numColumnsAtom);
   const filename = useFilename();
@@ -86,8 +89,8 @@ export const EditApp: React.FC<AppProps> = ({
 
   const { connection, reconnect, documentModel } = useNotebookSession({
     autoInstantiate: userConfig.runtime.auto_instantiate,
-    setCells: (cells, layout) => {
-      setCells(cells);
+    onKernelReady: (layout) => {
+      const cells = notebookCells(store.get(notebookAtom));
       const names = cells.map((cell) => cell.name);
       const codes = cells.map((cell) => cell.code);
       const configs = cells.map((cell) => cell.config);

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,