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

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

@@ -1,22 +1,29 @@
 /* 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";
 import { islandsHydratedAtom, islandsPendingInitialRunsAtom } from "../state";
-import type { IslandsKernelMessage } from "../worker/worker";
-
-function completedRun(sessionGeneration: number): IslandsKernelMessage {
-  return {
-    sessionGeneration,
-    message: JSON.stringify({
-      op: "completed-run",
-      data: { op: "completed-run", run_id: null },
-    }) as IslandsKernelMessage["message"],
-  };
-}
+import type { NotificationMessage } from "@/core/kernel/messages";
+
+type Consumer = (
+  notification: NotificationMessage,
+  sessionGeneration: number,
+) => void;
+
+const completedRun = {
+  op: "completed-run",
+  run_id: null,
+} as NotificationMessage;
 
 describe("initializeIslands", () => {
   afterEach(vi.restoreAllMocks);
@@ -26,12 +33,13 @@ describe("initializeIslands", () => {
     root.innerHTML =
       '<marimo-island data-app-id="app-1" data-reactive="true"></marimo-island>';
 
-    let consumeMessage: ((message: IslandsKernelMessage) => void) | undefined;
+    let consumeMessage: Consumer | undefined;
     const bridge = {
-      consumeMessages: (consumer: (message: IslandsKernelMessage) => void) => {
+      consumeNotifications: (consumer: Consumer) => {
         consumeMessage = consumer;
       },
       sendComponentValues: vi.fn(),
+      notebookBackend: fakeBackend(),
     } as unknown as IslandsPyodideBridge;
     const store = createStore();
     vi.spyOn(RuntimeState.INSTANCE, "start").mockImplementation(
@@ -48,14 +56,58 @@ describe("initializeIslands", () => {
     expect(store.get(islandsHydratedAtom)).toBe(false);
     store.set(islandsPendingInitialRunsAtom, new Set([1, 2]));
 
-    consumeMessage?.(completedRun(1));
+    consumeMessage?.(completedRun, 1);
     expect(store.get(islandsHydratedAtom)).toBe(false);
 
-    consumeMessage?.(completedRun(1));
-    consumeMessage?.(completedRun(0));
+    consumeMessage?.(completedRun, 1);
+    consumeMessage?.(completedRun, 0);
     expect(store.get(islandsHydratedAtom)).toBe(false);
 
-    consumeMessage?.(completedRun(2));
+    consumeMessage?.(completedRun, 2);
     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).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).toEqual(["c1", "c3", "k7f2"]);
+  });
+});

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

@@ -10,7 +10,6 @@ import {
 } from "@/__tests__/branded";
 import { notebookAtom } from "@/core/cells/cells";
 import { islandsPendingInitialRunsAtom } from "@/core/islands/state";
-import type { IslandsKernelMessage } from "../worker/worker";
 
 type Base64String = components["schemas"]["Base64String"];
 interface TestIslandApp {
@@ -89,9 +88,6 @@ vi.mock("@/core/wasm/rpc", () => ({
         startSession: mockStartSessionRequest,
         stopSession: mockStopSessionRequest,
       },
-      send: {
-        consumerReady: vi.fn(),
-      },
     },
     addMessageListener: vi.fn(
       (name: string, listener: (payload: never) => void) => {
@@ -177,17 +173,22 @@ describe("IslandsPyodideBridge", () => {
     });
   }
 
-  function sendKernelMessage(sessionGeneration: number) {
-    const listener = mockMessageListeners.get("kernelMessage");
+  function sendNotebookEvent(sessionGeneration: number, frame: unknown) {
+    const listener = mockMessageListeners.get("notebookEvent");
     if (!listener) {
-      throw new Error("Missing kernel message listener");
+      throw new Error("Missing notebook event listener");
     }
-    const message = {
-      message: "" as IslandsKernelMessage["message"],
-      sessionGeneration,
-    };
-    listener(message as never);
-    return message;
+    listener({ event: JSON.stringify(frame), sessionGeneration } as never);
+  }
+
+  /** Send a notification and return the call the consumer should receive. */
+  function sendKernelMessage(sessionGeneration: number) {
+    const notification = { op: "interrupted" };
+    sendNotebookEvent(sessionGeneration, {
+      event: "marimo",
+      data: { marimo: notification },
+    });
+    return [notification, sessionGeneration];
   }
 
   async function initializeSingleApp() {
@@ -300,7 +301,7 @@ describe("IslandsPyodideBridge", () => {
   describe("app lifecycle", () => {
     it("ignores messages after a session yields ownership", async () => {
       const consumeMessage = vi.fn();
-      bridge.consumeMessages(consumeMessage);
+      bridge.consumeNotifications(consumeMessage);
       mockSingleApp();
       await bridge.initializeApps();
 
@@ -319,10 +320,7 @@ describe("IslandsPyodideBridge", () => {
       sendKernelMessage(1);
       const currentMessage = sendKernelMessage(2);
 
-      expect(consumeMessage.mock.calls).toEqual([
-        [firstMessage],
-        [currentMessage],
-      ]);
+      expect(consumeMessage.mock.calls).toEqual([firstMessage, currentMessage]);
 
       finishReplacement();
       await replacement;
@@ -330,7 +328,7 @@ describe("IslandsPyodideBridge", () => {
 
     it("restores message ownership when replacement fails", async () => {
       const consumeMessage = vi.fn();
-      bridge.consumeMessages(consumeMessage);
+      bridge.consumeNotifications(consumeMessage);
       mockSingleApp();
       await bridge.initializeApps();
       mockReplaceSessionRequest.mockRejectedValueOnce(
@@ -344,12 +342,12 @@ describe("IslandsPyodideBridge", () => {
       const previousMessage = sendKernelMessage(1);
       sendKernelMessage(2);
 
-      expect(consumeMessage.mock.calls).toEqual([[previousMessage]]);
+      expect(consumeMessage.mock.calls).toEqual([previousMessage]);
     });
 
     it("ignores messages while their session is stopping", async () => {
       const consumeMessage = vi.fn();
-      bridge.consumeMessages(consumeMessage);
+      bridge.consumeNotifications(consumeMessage);
       mockSingleApp();
       await bridge.initializeApps();
       let finishStop!: () => void;
@@ -373,7 +371,7 @@ describe("IslandsPyodideBridge", () => {
 
     it("forwards messages from every app in a multi-app document", async () => {
       const consumeMessage = vi.fn();
-      bridge.consumeMessages(consumeMessage);
+      bridge.consumeNotifications(consumeMessage);
       mockParseMarimoIslandApps.mockReturnValue([app(), app("app-2")]);
       mockCreateMarimoFile
         .mockReturnValueOnce("generated app 1")
@@ -385,9 +383,58 @@ describe("IslandsPyodideBridge", () => {
       const secondMessage = sendKernelMessage(2);
       sendKernelMessage(3);
 
-      expect(consumeMessage.mock.calls).toEqual([
-        [firstMessage],
-        [secondMessage],
+      expect(consumeMessage.mock.calls).toEqual([firstMessage, secondMessage]);
+    });
+
+    it("delivers only the notebook stream's marimo notifications", async () => {
+      const consumeMessage = vi.fn();
+      bridge.consumeNotifications(consumeMessage);
+      await initializeSingleApp();
+      sendNotebookEvent(1, { event: "ready", id: "3", data: {} });
+      const notification = sendKernelMessage(1);
+
+      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: {} }],
       ]);
     });
 
@@ -723,8 +770,19 @@ describe("IslandsPyodideBridge", () => {
         functionName: "put_control_request",
         payload: {
           type: "execute-cells",
-          cellIds: ["cell-1", "cell-2"],
-          codes: ["print('hello')", "print('world')"],
+          request: undefined,
+          executionRequests: [
+            {
+              type: "execute-cell",
+              cellId: "cell-1",
+              code: "print('hello')",
+            },
+            {
+              type: "execute-cell",
+              cellId: "cell-2",
+              code: "print('world')",
+            },
+          ],
         },
         sessionGeneration: 1,
       });

package/src/core/islands/bootstrap.ts

@@ -9,11 +9,7 @@ import {
 } from "@/plugins/impl/anywidget/registry";
 import { initializePlugins } from "@/plugins/plugins";
 import { Functions } from "@/utils/functions";
-import {
-  safeExtractSetUIElementMessageBuffers,
-  type JsonString,
-} from "@/utils/json/base64";
-import { jsonParseWithSpecialChar } from "@/utils/json/json-parser";
+import { safeExtractSetUIElementMessageBuffers } from "@/utils/json/base64";
 import { Logger } from "@/utils/Logger";
 import { logNever } from "@/utils/assertNever";
 import {
@@ -21,19 +17,20 @@ 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,
   handleRemoveUIElements,
 } from "../kernel/handlers";
-import type {
-  NotificationMessage,
-  NotificationMessageType,
-} from "../kernel/messages";
+import type { NotificationMessage } from "../kernel/messages";
 import { queryParamHandlers } from "../kernel/queryParamHandlers";
 import { RuntimeState } from "../kernel/RuntimeState";
 import { initialModeAtom } from "../mode";
@@ -126,19 +123,26 @@ export async function initializeIslands(
     }
   });
 
-  // Wire up kernel message handling
-  // The wire format is {"op": "...", "data": {"op": "...", ...}} — both
-  // the envelope and the payload carry the op. The bridge types the message
-  // as NotificationPayload (just {data}), but the actual wire format
-  // includes the outer op too.
-  // Worker metadata identifies the session that emitted the message.
-  bridge.consumeMessages(({ message, sessionGeneration }) => {
-    handleMessage(
-      message as unknown as JsonString<IslandsNotificationMessage>,
-      actions,
-      store,
-      sessionGeneration,
-    );
+  // Each notification names the session that sent it.
+  bridge.consumeNotifications((notification, sessionGeneration) => {
+    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
@@ -154,27 +158,15 @@ export async function initializeIslands(
   RuntimeState.INSTANCE.start(bridge.sendComponentValues);
 }
 
-type IslandsNotificationMessage = {
-  [K in NotificationMessageType]: {
-    data: Extract<NotificationMessage, { op: K }>;
-    op: K;
-  };
-}[NotificationMessageType];
-
-/**
- * Handles a single message from the kernel.
- *
- * Wire format from Python: {"op": "<name>", "data": {"op": "<name>", ...}}
- */
+/** Handles a single notification from the kernel. */
 // oxlint-disable-next-line marimo/prefer-object-params
-function handleMessage(
-  message: JsonString<IslandsNotificationMessage>,
+function handleNotification(
+  msg: NotificationMessage,
   actions: NotebookActions,
   store: Store,
   sessionGeneration: number,
 ): void {
   try {
-    const msg = jsonParseWithSpecialChar(message);
     Logger.debug("Islands received message:", msg.op);
 
     switch (msg.op) {
@@ -223,11 +215,12 @@ function handleMessage(
         return;
 
       case "kernel-ready":
-        handleKernelReady(msg.data, {
+        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,
           setKernelState: Functions.NOOP,
           onError: Logger.error,
@@ -239,47 +232,47 @@ function handleMessage(
 
       case "send-ui-element-message":
         UI_ELEMENT_REGISTRY.broadcastMessage(
-          msg.data.ui_element,
-          msg.data.message,
-          safeExtractSetUIElementMessageBuffers(msg.data),
+          msg.ui_element,
+          msg.message,
+          safeExtractSetUIElementMessageBuffers(msg),
         );
         return;
 
       case "remove-ui-elements":
-        handleRemoveUIElements(msg.data);
+        handleRemoveUIElements(msg);
         return;
 
       case "function-call-result":
-        FUNCTIONS_REGISTRY.resolve(msg.data.function_call_id, msg.data);
+        FUNCTIONS_REGISTRY.resolve(msg.function_call_id, msg);
         return;
 
       case "cell-op":
-        handleCellNotificationeration(msg.data, actions.handleCellMessage);
+        handleCellNotificationeration(msg, actions.handleCellMessage);
         return;
 
       case "alert":
         toast({
-          title: msg.data.title,
-          description: renderHTML({ html: msg.data.description }),
-          variant: msg.data.variant,
+          title: msg.title,
+          description: renderHTML({ html: msg.description }),
+          variant: msg.variant,
         });
         return;
 
       case "query-params-append":
-        queryParamHandlers.append(msg.data);
+        queryParamHandlers.append(msg);
         return;
       case "query-params-set":
-        queryParamHandlers.set(msg.data);
+        queryParamHandlers.set(msg);
         return;
       case "query-params-delete":
-        queryParamHandlers.delete(msg.data);
+        queryParamHandlers.delete(msg);
         return;
       case "query-params-clear":
         queryParamHandlers.clear();
         return;
 
       case "model-lifecycle":
-        handleWidgetMessage(WIDGET_REGISTRY, msg.data);
+        handleWidgetMessage(WIDGET_REGISTRY, msg);
         return;
 
       case "consumer-capabilities":