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

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

@@ -0,0 +1,166 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+import { NotebookEvents, StreamStopped } from "@marimo-team/client";
+import { afterEach, assert, describe, expect, it, vi } from "vitest";
+import { Deferred } from "@/utils/Deferred";
+import { createWasmNotebookBackend, type WasmNotebookBridge } from "../wasm";
+import { fixtureEvents } from "../../notebook/replica/__tests__/fixtures";
+
+const streams: NotebookEvents[] = [];
+afterEach(async () => {
+  for (const stream of streams.splice(0)) {
+    stream.close();
+    await stream.closed.catch(() => undefined);
+  }
+});
+
+function worker() {
+  const frames = fixtureEvents("notebook-events.sse");
+  const document = frames.find((event) => event.event === "document")?.data
+    .document;
+  assert(document);
+  const listeners = new Set<(event: unknown) => void>();
+  const reply = { ok: true, status: 200, data: { document, cells: {} } };
+  const wire = { status: 200, body: reply.data };
+  const bridge = {
+    subscribe: vi.fn(
+      async (_cursor: string | null, emit: (event: unknown) => void) => {
+        listeners.add(emit);
+        frames.forEach(emit);
+        return () => {
+          listeners.delete(emit);
+        };
+      },
+    ),
+    call: vi.fn<WasmNotebookBridge["call"]>().mockResolvedValue(wire),
+  };
+  return { bridge, listeners, reply, wire };
+}
+
+async function follow(bridge: WasmNotebookBridge) {
+  const stream = new NotebookEvents(
+    createWasmNotebookBackend(bridge).transport,
+  );
+  streams.push(stream);
+  const ready = new Deferred<void>();
+  stream.subscribe((event) => {
+    if (event.event === "ready") {
+      ready.resolve();
+    }
+  }, ready.reject);
+  await ready.promise;
+  return stream;
+}
+
+describe("WASM notebook backend", () => {
+  it("follows independent readers and releases each subscription on close", async () => {
+    const { bridge, listeners } = worker();
+    const first = await follow(bridge);
+    const second = await follow(bridge);
+    expect(listeners.size).toBe(2);
+    expect(second.status).toBe("connected");
+    first.close();
+    await first.closed;
+    expect(listeners.size).toBe(1);
+    second.close();
+    await second.closed;
+    expect(listeners.size).toBe(0);
+  });
+
+  it("closes during startup and releases the late subscription", async () => {
+    const { bridge } = worker();
+    const pending = new Deferred<() => void>();
+    bridge.subscribe.mockReturnValue(pending.promise);
+    const controller = new AbortController();
+    const done = createWasmNotebookBackend(bridge).transport(
+      { cursor: "previous-4", signal: controller.signal },
+      { open: vi.fn(), message: vi.fn() },
+    );
+    expect(bridge.subscribe).toHaveBeenCalledWith(
+      "previous-4",
+      expect.any(Function),
+      controller.signal,
+    );
+    controller.abort();
+    await done;
+    const close = vi.fn();
+    pending.resolve(close);
+    await vi.waitFor(() => expect(close).toHaveBeenCalledOnce());
+  });
+
+  it("stops on worker initialization failure", async () => {
+    const { bridge } = worker();
+    bridge.subscribe.mockRejectedValue(new Error("Initialization failed"));
+    await expect(
+      createWasmNotebookBackend(bridge).transport(
+        { cursor: null, signal: new AbortController().signal },
+        { open: vi.fn(), message: vi.fn() },
+      ),
+    ).rejects.toBeInstanceOf(StreamStopped);
+  });
+
+  it.each(["edit", "save"] as const)(
+    "sends %s calls with the retry key and validates replies",
+    async (operation) => {
+      const { bridge, reply } = worker();
+      const notebook = createWasmNotebookBackend(bridge).notebook("n1");
+      const command =
+        operation === "edit"
+          ? notebook.edit.bind(notebook, [])
+          : notebook.save.bind(notebook);
+      expect(await command({ idempotencyKey: `${operation}-7` })).toEqual(
+        reply,
+      );
+      expect(bridge.call).toHaveBeenCalledWith({
+        operation: `notebook.${operation}`,
+        params: { notebook_id: "n1" },
+        headers: { "idempotency-key": `${operation}-7` },
+        body: operation === "edit" ? { changes: [] } : {},
+      });
+      const problem = {
+        type: "tag:marimo.io,2026:host/forbidden",
+        title: "Forbidden",
+        status: 403,
+      };
+      bridge.call.mockResolvedValue({ status: 403, body: problem });
+      expect(await command()).toEqual({
+        ok: false,
+        status: 403,
+        problem,
+      });
+      bridge.call.mockResolvedValue({ status: 200, body: {} });
+      await expect(command()).rejects.toThrow(/Invalid input/);
+    },
+  );
+
+  it.each(["edit", "save"] as const)(
+    "cancels waiting for an accepted %s and permits retry with the same key",
+    async (operation) => {
+      const { bridge, reply, wire } = worker();
+      const pending = new Deferred<unknown>();
+      bridge.call.mockReturnValue(pending.promise);
+      const notebook = createWasmNotebookBackend(bridge).notebook("n1");
+      const command =
+        operation === "edit"
+          ? notebook.edit.bind(notebook, [])
+          : notebook.save.bind(notebook);
+      const controller = new AbortController();
+      const pendingCommand = command({
+        signal: controller.signal,
+        idempotencyKey: `${operation}-8`,
+      });
+      await vi.waitFor(() => expect(bridge.call).toHaveBeenCalledOnce());
+      controller.abort();
+      await expect(pendingCommand).rejects.toBe(controller.signal.reason);
+      pending.resolve(wire);
+      expect(await command({ idempotencyKey: `${operation}-8` })).toEqual(
+        reply,
+      );
+      expect(bridge.call.mock.calls[0]).toEqual(bridge.call.mock.calls[1]);
+      await expect(command({ signal: controller.signal })).rejects.toBe(
+        controller.signal.reason,
+      );
+      expect(bridge.call).toHaveBeenCalledTimes(2);
+    },
+  );
+});

package/src/core/{websocket → transports}/connection-state.ts

@@ -1,23 +1,16 @@
 /* Copyright 2026 Marimo. All rights reserved. */
 
-import type { CloseDecision } from "../host/connection";
+import {
+  type CloseDecision,
+  ConnectionState,
+  ConnectionClosedReason,
+} from "../notebook/replica";
 import { Logger } from "@/utils/Logger";
-import { TRANSPORT_EXHAUSTED_REASON } from "./transports/ws";
-import { ConnectionState, WebSocketClosedReason } from "./types";
 
 // Keep server reasons in sync with ws_endpoint.py and ws/*.py.
 // SSE carries the same reasons in close events. Unclassified reasons retry.
 export function classifyCloseEvent(event: { reason?: string }): CloseDecision {
   switch (event.reason) {
-    case TRANSPORT_EXHAUSTED_REASON:
-      return {
-        kind: "gave-up",
-        status: {
-          state: ConnectionState.CLOSED,
-          code: WebSocketClosedReason.KERNEL_DISCONNECTED,
-          reason: "kernel not found",
-        },
-      };
     case "MARIMO_NO_FILE_KEY":
     case "MARIMO_NO_SESSION_ID":
     case "MARIMO_NO_SESSION":
@@ -26,7 +19,7 @@ export function classifyCloseEvent(event: { reason?: string }): CloseDecision {
         kind: "terminal",
         status: {
           state: ConnectionState.CLOSED,
-          code: WebSocketClosedReason.KERNEL_DISCONNECTED,
+          code: ConnectionClosedReason.KERNEL_DISCONNECTED,
           reason: "kernel not found",
         },
       };
@@ -35,7 +28,7 @@ export function classifyCloseEvent(event: { reason?: string }): CloseDecision {
         kind: "terminal",
         status: {
           state: ConnectionState.CLOSED,
-          code: WebSocketClosedReason.KERNEL_DISCONNECTED,
+          code: ConnectionClosedReason.KERNEL_DISCONNECTED,
           reason: "not authorized",
         },
       };
@@ -44,7 +37,7 @@ export function classifyCloseEvent(event: { reason?: string }): CloseDecision {
         kind: "terminal",
         status: {
           state: ConnectionState.CLOSED,
-          code: WebSocketClosedReason.KERNEL_DISCONNECTED,
+          code: ConnectionClosedReason.KERNEL_DISCONNECTED,
           reason: "kiosk mode is not available for this session",
         },
       };
@@ -53,7 +46,7 @@ export function classifyCloseEvent(event: { reason?: string }): CloseDecision {
         kind: "terminal",
         status: {
           state: ConnectionState.CLOSED,
-          code: WebSocketClosedReason.KERNEL_STARTUP_ERROR,
+          code: ConnectionClosedReason.KERNEL_STARTUP_ERROR,
           reason: "Failed to start kernel sandbox",
         },
       };

package/src/core/transports/server.ts

@@ -0,0 +1,152 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+import {
+  type NotebookTransport,
+  sse,
+  StreamStopped,
+  type StreamUrl,
+} from "@marimo-team/client";
+import { jsonParseWithSpecialChar } from "@/utils/json/json-parser";
+import { classifyCloseEvent } from "./connection-state";
+
+/** The server can be slow to accept when many notebooks are open. */
+const CONNECT_MS = 10_000;
+
+/**
+ * The error that ends a stream on this close, or none to retry.
+ *
+ * `cause` is the connection status the editor shows.
+ */
+function stopOn(close: { code?: number; reason?: string }) {
+  const decision = classifyCloseEvent(close);
+  if (decision.kind === "retry") {
+    return undefined;
+  }
+  const reason = "reason" in decision.status ? decision.status.reason : "";
+  return new StreamStopped(reason, { cause: decision.status });
+}
+
+/** Read a URL provider, stopping the stream when it gives up. */
+async function resolve(url: StreamUrl): Promise<URL | undefined> {
+  try {
+    return new URL(typeof url === "function" ? await url() : url);
+  } catch (error) {
+    if (error instanceof StreamStopped) {
+      throw error;
+    }
+    return undefined;
+  }
+}
+
+/** Send each request with the runtime's headers, read before each request. */
+export function withHeaders(
+  headers: () => Record<string, string>,
+): typeof globalThis.fetch {
+  return (input, init) => {
+    const sent = new Headers(init?.headers);
+    for (const [name, value] of Object.entries(headers())) {
+      sent.set(name, value);
+    }
+    return globalThis.fetch(input, { ...init, headers: sent });
+  };
+}
+
+/** A resume cursor and replica belong to the session that produced them. */
+function withSessionIdentity(
+  url: StreamUrl,
+  connect: (target: URL) => NotebookTransport,
+): NotebookTransport {
+  let previous: string | undefined;
+  return async (attempt, handlers) => {
+    const target = await resolve(url);
+    if (target === undefined || attempt.signal.aborted) {
+      return;
+    }
+    target.searchParams.delete("cursor");
+    const identity = new URL(target);
+    identity.searchParams.delete("access_token");
+    identity.searchParams.sort();
+    identity.hash = "";
+    const key = identity.toString();
+    let { cursor } = attempt;
+    if (previous !== undefined && previous !== key) {
+      // Reset both the stream's cursor and its subscribers' state.
+      handlers.message({ event: "reset", data: {} });
+      cursor = null;
+    }
+    previous = key;
+    if (!attempt.signal.aborted) {
+      return connect(target)({ ...attempt, cursor }, handlers);
+    }
+  };
+}
+
+/**
+ * Follow a marimo session over its WebSocket route.
+ *
+ * Each attempt opens one socket with the cursor as a query parameter. The
+ * editor's close policy decides which closes end the stream.
+ *
+ * TODO: Notice a connection that died without closing, such as after a laptop
+ * sleeps. `/sse` sends a heartbeat that its silence limit relies on, and
+ * `/ws` drops it. Sending it on `/ws` would let this transport close a
+ * silent socket.
+ */
+export function websocket(url: StreamUrl): NotebookTransport {
+  return withSessionIdentity(
+    url,
+    (target) =>
+      async ({ cursor, signal }, { open, message }) => {
+        if (cursor !== null) {
+          target.searchParams.set("cursor", cursor);
+        }
+        await new Promise<void>((settle, reject) => {
+          const socket = new WebSocket(target);
+          const timeout = setTimeout(() => socket.close(), CONNECT_MS);
+          const abort = () => socket.close();
+          signal.addEventListener("abort", abort, { once: true });
+          socket.addEventListener("open", () => {
+            clearTimeout(timeout);
+            open();
+          });
+          socket.addEventListener("message", ({ data }) =>
+            // marimo's parser reads `NaN`. A message it cannot read arrives as an
+            // empty object, which `NotebookEvents` rejects as not an event.
+            message(jsonParseWithSpecialChar(String(data))),
+          );
+          socket.addEventListener("close", (event) => {
+            clearTimeout(timeout);
+            signal.removeEventListener("abort", abort);
+            const stopped = signal.aborted ? undefined : stopOn(event);
+            if (stopped === undefined) {
+              settle();
+            } else {
+              reject(stopped);
+            }
+          });
+        });
+      },
+  );
+}
+
+/**
+ * Follow a marimo session over its server-sent events route.
+ *
+ * The session's `close` frames carry the same reasons as WebSocket closes.
+ * marimo's parser decodes tagged bigints and non-finite numbers in every frame.
+ */
+export function sessionSse(
+  url: StreamUrl,
+  headers: () => Record<string, string>,
+): NotebookTransport {
+  return withSessionIdentity(url, (target) =>
+    sse(target, {
+      parseData: jsonParseWithSpecialChar,
+      fetch: withHeaders(headers),
+      stop: (frame) =>
+        frame.event === "close"
+          ? stopOn(frame.data as { code?: number; reason?: string })
+          : undefined,
+    }),
+  );
+}

package/src/core/transports/static.ts

@@ -0,0 +1,84 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+/** A backend that serves a document snapshot without a host. */
+
+import type { NotebookClient, NotebookTransport } 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[];
+}
+
+/**
+ * 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[] = [
+    {
+      event: "notebook",
+      data: { notebook: { ...snapshot.notebook, operations: [] } },
+    },
+    ...snapshot.cells.map((cell): Protocol.CellEvent => ({
+      event: "cell",
+      data: { cell },
+    })),
+    { event: "document", data: { document: snapshot.document } },
+    { 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,
+      };
+    },
+  };
+}
+
+/** A connection that opens and delivers nothing, for static notebooks. */
+export const idle: NotebookTransport = ({ signal }, { open }) =>
+  new Promise((settle) => {
+    open();
+    signal.addEventListener("abort", () => settle(), { once: true });
+  });

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 }),
+  };
+}