@marimo-team/islands 0.25.3-dev20 → 0.25.3-dev201

package/src/core/transports/static.ts

@@ -0,0 +1,166 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+/** A backend that serves a document snapshot without a host. */
+
+import type { NotebookClient } from "@marimo-team/client";
+import type * as Protocol from "@marimo-team/client/protocol";
+import type { NotebookBackend } from "../notebook/replica";
+
+/** A document snapshot with its own identity and revision metadata. */
+export interface StaticNotebookSnapshot {
+  readonly notebook: Omit<Protocol.Notebook, "operations">;
+  readonly document: Protocol.Document;
+  readonly cells: readonly Protocol.Cell[];
+}
+
+/** The events that describe a snapshot, in the order a host sends them. */
+function snapshotFrames(snapshot: StaticNotebookSnapshot): Protocol.Event[] {
+  return [
+    {
+      event: "notebook",
+      data: { notebook: { ...snapshot.notebook, operations: [] } },
+    },
+    ...snapshot.cells.map((cell): Protocol.CellEvent => ({
+      event: "cell",
+      data: { cell },
+    })),
+    { event: "document", data: { document: snapshot.document } },
+  ];
+}
+
+/** A snapshot of a notebook with no cells. */
+export function emptySnapshot(notebookId: string): StaticNotebookSnapshot {
+  return {
+    notebook: {
+      id: notebookId,
+      project_id: notebookId,
+      path: null,
+      name: "",
+      updated_at: null,
+    },
+    document: {
+      notebook_id: notebookId,
+      order: [],
+      app_config: {},
+      etag: `${notebookId}:document`,
+      saved_etag: `${notebookId}:document`,
+      order_etag: `${notebookId}:order`,
+      app_config_etag: `${notebookId}:app_config`,
+      version: 1,
+    },
+    cells: [],
+  };
+}
+
+/**
+ * Show a snapshot until `live` answers.
+ *
+ * The page carries the notebook it was served with, so cells appear before
+ * the host connects. The snapshot ends in a `ready` of its own, like a
+ * static notebook's, so the replica publishes it. At the host's first sign
+ * of life a `reset` discards it, and the host's snapshot takes over. A
+ * connection that resumes from a host cursor skips the snapshot; one that
+ * resumes from the snapshot's cursor starts the host connection fresh,
+ * since the host never issued it.
+ */
+export function startWithSnapshot(
+  snapshot: StaticNotebookSnapshot,
+  live: NotebookBackend,
+): NotebookBackend {
+  const cursor = `${snapshot.notebook.id}:snapshot`;
+  return {
+    ...live,
+    transport: (attempt, handlers) => {
+      if (attempt.cursor !== null && attempt.cursor !== cursor) {
+        return live.transport(attempt, handlers);
+      }
+      if (attempt.cursor === null) {
+        for (const frame of snapshotFrames(snapshot)) {
+          if (attempt.signal.aborted) {
+            break;
+          }
+          handlers.message(frame);
+        }
+        if (attempt.signal.aborted) {
+          return Promise.resolve();
+        }
+        handlers.message({ event: "ready", id: cursor, data: {} });
+      }
+      let shown = true;
+      const discard = () => {
+        if (shown) {
+          shown = false;
+          handlers.message({ event: "reset", data: {} });
+        }
+      };
+      return live.transport(
+        { ...attempt, cursor: null },
+        {
+          open: () => {
+            discard();
+            handlers.open();
+          },
+          message: (frame) => {
+            discard();
+            handlers.message(frame);
+          },
+        },
+      );
+    },
+  };
+}
+
+/**
+ * Serve a document without a server, runtime, or mutation operations.
+ * Callers and readers must treat the snapshot records as immutable.
+ */
+export function createStaticNotebookBackend(
+  snapshot: StaticNotebookSnapshot,
+): NotebookBackend {
+  const frames: Protocol.Event[] = [
+    ...snapshotFrames(snapshot),
+    { event: "ready", id: "static", data: {} },
+  ];
+
+  return {
+    transport: ({ cursor, signal }, { open, message }) =>
+      new Promise((resolve) => {
+        if (signal.aborted) {
+          resolve();
+          return;
+        }
+        // Finishing would ask NotebookEvents to reconnect. A snapshot stays open
+        // until its reader closes, even though there are no more updates.
+        signal.addEventListener("abort", () => resolve(), { once: true });
+        open();
+        if (cursor !== null && !signal.aborted) {
+          message({ event: "reset", data: {} });
+        }
+        for (const frame of frames) {
+          if (signal.aborted) {
+            return;
+          }
+          message(frame);
+        }
+      }),
+    notebook: () => {
+      const unsupported: NotebookClient["save"] = async (options) => {
+        options?.signal?.throwIfAborted();
+        return {
+          ok: false,
+          status: 501,
+          problem: {
+            type: "tag:marimo.io,2026:host/unsupported-operation",
+            title: "Unsupported operation",
+            status: 501,
+            detail: "Static notebooks do not support mutations",
+          },
+        };
+      };
+      return {
+        edit: (_changes, options) => unsupported(options),
+        save: unsupported,
+      };
+    },
+  };
+}

package/src/core/transports/wasm.ts

@@ -0,0 +1,97 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+/** A backend for the notebook running in the WASM worker. */
+
+import {
+  NotebookClient,
+  StreamStopped,
+  type unstable_NotebookCall,
+} from "@marimo-team/client";
+import type { NotebookBackend } from "../notebook/replica";
+
+/** The worker owns the notebook; this interface only carries commands and events. */
+export interface WasmNotebookBridge {
+  subscribe(
+    cursor: string | null,
+    emit: (event: unknown) => void,
+    signal: AbortSignal,
+  ): Promise<() => void>;
+  /** Deliver a call to the worker and resolve with its status and body. */
+  call(call: unstable_NotebookCall): Promise<unknown>;
+}
+
+/** Stop waiting on cancellation; a command already sent may still commit. */
+function abortable<T>(
+  send: () => Promise<T>,
+  signal?: AbortSignal,
+): Promise<T> {
+  if (!signal) {
+    return send();
+  }
+  signal.throwIfAborted();
+  return new Promise((resolve, reject) => {
+    const abort = () => reject(signal.reason);
+    signal.addEventListener("abort", abort, { once: true });
+    send()
+      .then(resolve, reject)
+      .finally(() => signal.removeEventListener("abort", abort));
+  });
+}
+
+export function createWasmNotebookBackend(
+  bridge: WasmNotebookBridge,
+): NotebookBackend {
+  return {
+    transport: async ({ cursor, signal }, { open, message }) => {
+      if (signal.aborted) {
+        return;
+      }
+      let unsubscribe: (() => void) | undefined;
+      let finish!: () => void;
+      const closed = new Promise<void>((resolve) => {
+        finish = resolve;
+      });
+      const abort = () => {
+        unsubscribe?.();
+        finish();
+      };
+      signal.addEventListener("abort", abort, { once: true });
+      try {
+        open();
+        // Keep the pending subscription so a close during startup still releases it.
+        await abortable(
+          () =>
+            bridge
+              .subscribe(
+                cursor,
+                (event) => {
+                  if (!signal.aborted) {
+                    message(event);
+                  }
+                },
+                signal,
+              )
+              .then((close) => {
+                unsubscribe = close;
+                if (signal.aborted) {
+                  close();
+                }
+              }),
+          signal,
+        );
+        await closed;
+      } catch (error) {
+        if (!signal.aborted) {
+          unsubscribe?.();
+          throw new StreamStopped("Could not follow the WASM notebook", {
+            cause: error,
+          });
+        }
+      } finally {
+        signal.removeEventListener("abort", abort);
+      }
+    },
+    notebook: (notebookId) =>
+      NotebookClient.unstable_over((call) => bridge.call(call), { notebookId }),
+  };
+}

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

@@ -8,16 +8,36 @@ const {
   mockBridge,
   mockNotebookReadFile,
   mockReadNotebook,
-  mockSaveNotebook,
+  mockWriteNotebook,
+  mockNotebookCall,
+  mockSetNotebookLayout,
+  mockFlush,
+  parentRequests,
   mockBootstrap,
   rpcListeners,
+  mockNotebookSubscribe,
+  mockNotebookUnsubscribe,
 } = vi.hoisted(() => ({
   mockBridge: vi.fn(),
   mockNotebookReadFile: vi.fn(),
   mockReadNotebook: vi.fn(),
-  mockSaveNotebook: vi.fn(),
+  mockWriteNotebook: vi.fn(),
+  mockNotebookCall: vi.fn(),
+  mockSetNotebookLayout: vi.fn(),
+  mockFlush: vi.fn(),
+  parentRequests: {} as {
+    persistNotebook?: (snapshot: {
+      filename: string;
+      contents: string;
+    }) => Promise<void>;
+  },
   mockBootstrap: vi.fn(),
-  rpcListeners: {} as Record<string, () => void>,
+  rpcListeners: {} as Record<
+    string,
+    (message?: { id: string; event: string }) => void
+  >,
+  mockNotebookSubscribe: vi.fn(),
+  mockNotebookUnsubscribe: vi.fn(),
 }));
 
 // Mock browser globals before any imports
@@ -41,22 +61,31 @@ class MockURL extends URL {
 vi.stubGlobal("URL", MockURL);
 
 vi.mock("@/core/wasm/rpc", () => ({
-  getWorkerRPC: () => ({
-    send: { bootstrap: mockBootstrap },
-    proxy: {
-      request: {
-        bridge: mockBridge,
-        startSession: vi.fn(),
-        readFile: vi.fn(),
-        readNotebook: mockReadNotebook,
-        saveNotebook: mockSaveNotebook,
+  getWorkerRPC: (_worker: Worker, handlers?: typeof parentRequests) => {
+    Object.assign(parentRequests, handlers);
+    return {
+      send: { bootstrap: mockBootstrap },
+      proxy: {
+        request: {
+          bridge: mockBridge,
+          startSession: vi.fn(),
+          readFile: vi.fn(),
+          readNotebook: mockReadNotebook,
+          writeNotebook: mockWriteNotebook,
+          notebookSubscribe: mockNotebookSubscribe,
+          notebookUnsubscribe: mockNotebookUnsubscribe,
+          notebookCall: mockNotebookCall,
+          setNotebookLayout: mockSetNotebookLayout,
+        },
       },
-      send: { consumerReady: vi.fn() },
-    },
-    addMessageListener: (event: string, cb: () => void) => {
-      rpcListeners[event] = cb;
-    },
-  }),
+      addMessageListener: (
+        event: string,
+        cb: (typeof rpcListeners)[string],
+      ) => {
+        rpcListeners[event] = cb;
+      },
+    };
+  },
 }));
 
 vi.mock("@/core/meta/globals", () => ({
@@ -72,14 +101,25 @@ vi.mock("@/core/wasm/store", () => ({
   notebookFileStore: { readFile: mockNotebookReadFile, saveFile: vi.fn() },
 }));
 
+vi.mock("@/core/cells/document", async (importOriginal) => ({
+  ...(await importOriginal<typeof import("@/core/cells/document")>()),
+  flushDocument: mockFlush,
+}));
+
 // Import after all mocks are set up
 import { store } from "@/core/state/jotai";
 import { initialModeAtom } from "@/core/mode";
+import { fixtureEvents } from "@/core/notebook/replica/__tests__/fixtures";
+import { createNotebookReplica } from "@/core/notebook/replica";
+import { activeNotebookReplicaAtom } from "@/core/session/active-notebook";
+import { createEditorBackend } from "@/core/session/backend";
+import { fallbackFileStore, notebookFileStore } from "../store";
 import { getWasmWorkerName, PyodideBridge } from "../bridge";
 
-// Access INSTANCE once at module level so the constructor runs (and
-// addMessageListener populates rpcListeners) before any test executes.
-void PyodideBridge.INSTANCE;
+beforeEach(() => {
+  Reflect.deleteProperty(window, "_marimo_private_PyodideBridge");
+  void PyodideBridge.INSTANCE;
+});
 
 describe("PyodideBridge.readCode", () => {
   beforeEach(() => {
@@ -161,6 +201,7 @@ describe("PyodideBridge.exportAsScript", () => {
 });
 
 describe("PyodideBridge.sendSave", () => {
+  let disconnect: () => void;
   const request = {
     cellIds: [cellId("cell-1")],
     codes: ['value = "NEW"'],
@@ -171,51 +212,174 @@ describe("PyodideBridge.sendSave", () => {
     persist: true,
   };
 
-  beforeEach(() => {
+  beforeEach(async () => {
     vi.clearAllMocks();
     store.set(initialModeAtom, "edit");
     rpcListeners.initialized();
-    mockSaveNotebook.mockResolvedValue(null);
-    mockReadNotebook.mockResolvedValue(
-      "import marimo\napp = marimo.App()\n# NEW",
+    mockFlush.mockResolvedValue(undefined);
+    mockWriteNotebook.mockResolvedValue(undefined);
+    mockNotebookUnsubscribe.mockResolvedValue(undefined);
+    const document = fixtureEvents("notebook-events.sse").find(
+      (event) => event.event === "document",
+    )?.data.document;
+    mockNotebookCall.mockResolvedValue({
+      status: 200,
+      body: { document, cells: {} },
+    });
+    mockNotebookSubscribe.mockImplementation(async ({ id }) => {
+      for (const event of fixtureEvents("notebook-events.sse")) {
+        rpcListeners.notebookEvent({ id, event: JSON.stringify(event) });
+      }
+    });
+    const replica = createNotebookReplica(store);
+    store.set(activeNotebookReplicaAtom, replica);
+    const backend = createEditorBackend({
+      static: false,
+      snapshot: null,
+      transportType: "sse",
+      url: () => "http://localhost/",
+      headers: () => ({}),
+      notebookUrl: (id) => new URL(`http://localhost/notebooks/${id}`),
+      waitUntilReady: async () => undefined,
+      checkHealth: async () => true,
+    });
+    disconnect = replica.connect(backend, { onNotification: vi.fn() });
+    await vi.waitFor(() =>
+      expect(store.get(replica.notebookAtom)?.id).toBe("n1"),
     );
   });
 
   afterEach(() => {
+    disconnect();
+    store.set(activeNotebookReplicaAtom, null);
     store.set(initialModeAtom, undefined);
   });
 
-  it("waits for the session save before generating exports", async () => {
-    const mainSave = new Deferred<void>();
-    mockSaveNotebook
-      .mockResolvedValueOnce(null)
-      .mockReturnValueOnce(mainSave.promise);
+  it("saves through the connected backend using the stream's notebook ID", async () => {
+    const bridge = PyodideBridge.INSTANCE;
+    const createBackend = vi.spyOn(bridge, "createNotebookBackend");
+    await bridge.sendSave(request);
+    expect(createBackend).not.toHaveBeenCalled();
+    expect(mockNotebookCall).toHaveBeenCalledWith(
+      expect.objectContaining({
+        operation: "notebook.save",
+        params: { notebook_id: "n1" },
+        body: {},
+      }),
+    );
+  });
 
-    await PyodideBridge.INSTANCE.sendSave(request);
-    expect(mockSaveNotebook).toHaveBeenCalledTimes(2);
-    expect(mockSaveNotebook).toHaveBeenNthCalledWith(1, request);
-    expect(mockSaveNotebook).toHaveBeenNthCalledWith(2, request);
+  it.each(["unmounted", "not ready", "disconnected"])(
+    "does not save when the notebook is %s",
+    async (state) => {
+      if (state === "unmounted") {
+        store.set(activeNotebookReplicaAtom, null);
+      } else if (state === "not ready") {
+        store.set(activeNotebookReplicaAtom, createNotebookReplica(store));
+      } else {
+        disconnect();
+      }
+      await expect(PyodideBridge.INSTANCE.sendSave(request)).rejects.toThrow(
+        "The notebook is not connected",
+      );
+      expect(mockNotebookCall).not.toHaveBeenCalled();
+    },
+  );
 
-    const markdown = PyodideBridge.INSTANCE.exportAsMarkdown({
-      download: false,
-    });
-    const script = PyodideBridge.INSTANCE.exportAsScript({
-      download: false,
-    });
+  it("flushes edits before asking the host to save without an editor snapshot", async () => {
+    const flushed = new Deferred<void>();
+    mockFlush.mockReturnValueOnce(flushed.promise);
+    const saving = PyodideBridge.INSTANCE.sendSave(request);
     await Promise.resolve();
-    expect(mockBridge).not.toHaveBeenCalled();
+    expect(mockNotebookCall).not.toHaveBeenCalled();
+    flushed.resolve();
+    await saving;
+    expect(mockSetNotebookLayout).toHaveBeenCalledWith(null);
+    expect(mockNotebookCall).toHaveBeenCalledWith(
+      expect.objectContaining({ operation: "notebook.save", body: {} }),
+    );
+    expect(mockWriteNotebook).not.toHaveBeenCalled();
+  });
 
-    mainSave.resolve();
-    await Promise.all([markdown, script]);
-    expect(mockBridge).toHaveBeenCalledWith({
-      functionName: "export_markdown",
-      payload: { download: false },
-    });
-    expect(mockBridge).toHaveBeenCalledWith({
-      functionName: "export_script",
-      payload: { download: false },
+  it("does not save when flushing fails", async () => {
+    mockFlush.mockRejectedValueOnce(new Error("unsynchronized edit"));
+    await expect(PyodideBridge.INSTANCE.sendSave(request)).rejects.toThrow(
+      "unsynchronized edit",
+    );
+    expect(mockNotebookCall).not.toHaveBeenCalled();
+  });
+
+  it("reports host persistence failures", async () => {
+    mockNotebookCall.mockResolvedValueOnce({
+      status: 409,
+      body: {
+        type: "tag:marimo.io,2026:host/conflict",
+        title: "Conflict",
+        status: 409,
+        detail: "storage full",
+      },
     });
+    await expect(PyodideBridge.INSTANCE.sendSave(request)).rejects.toThrow(
+      "storage full",
+    );
   });
+
+  it("persists exactly the host's content before acknowledging storage", async () => {
+    const durable = new Deferred<void>();
+    const snapshot = { filename: "notebook.py", contents: "# host revision" };
+    mockWriteNotebook.mockReturnValueOnce(durable.promise);
+    const writing = parentRequests.persistNotebook!(snapshot);
+    await Promise.resolve();
+    expect(mockWriteNotebook).toHaveBeenCalledWith(snapshot);
+    expect(notebookFileStore.saveFile).not.toHaveBeenCalled();
+    durable.resolve();
+    await writing;
+    expect(notebookFileStore.saveFile).toHaveBeenCalledWith(snapshot.contents);
+    expect(fallbackFileStore.saveFile).toHaveBeenCalledWith(snapshot.contents);
+  });
+
+  it("propagates storage failure without updating the URL", async () => {
+    mockWriteNotebook.mockRejectedValueOnce(new Error("IndexedDB failed"));
+    await expect(
+      parentRequests.persistNotebook!({
+        filename: "notebook.py",
+        contents: "# host",
+      }),
+    ).rejects.toThrow("IndexedDB failed");
+    expect(notebookFileStore.saveFile).not.toHaveBeenCalled();
+  });
+});
+
+it("cancels a pending notebook subscription before its RPC response", async () => {
+  vi.clearAllMocks();
+  store.set(initialModeAtom, "read");
+  rpcListeners.initialized();
+  const pending = new Deferred<void>();
+  mockNotebookSubscribe.mockReturnValue(pending.promise);
+  mockNotebookUnsubscribe.mockResolvedValue(undefined);
+  const controller = new AbortController();
+  try {
+    const backend = await PyodideBridge.INSTANCE.createNotebookBackend();
+    const following = backend.transport(
+      { cursor: null, signal: controller.signal },
+      { open: vi.fn(), message: vi.fn() },
+    );
+    await vi.waitFor(() =>
+      expect(mockNotebookSubscribe).toHaveBeenCalledOnce(),
+    );
+    const [{ id }] = mockNotebookSubscribe.mock.calls[0];
+    controller.abort();
+    await following;
+    expect(mockNotebookUnsubscribe.mock.calls).toEqual([[id]]);
+
+    pending.resolve();
+    await new Promise((resolve) => setTimeout(resolve, 0));
+    expect(mockNotebookUnsubscribe.mock.calls).toEqual([[id]]);
+  } finally {
+    controller.abort();
+    pending.resolve();
+    store.set(initialModeAtom, undefined);
+  }
 });
 
 describe("getWasmWorkerName", () => {