@marimo-team/islands 0.25.3-dev23 → 0.25.3-dev230

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

@@ -0,0 +1,306 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+import { createNotebookReducer, NotebookEvents } from "@marimo-team/client";
+import * as Protocol from "@marimo-team/client/protocol";
+import { afterEach, assert, describe, expect, it, vi } from "vitest";
+import type { NotebookBackend } from "../../notebook/replica";
+import { document as documentResource } from "../../notebook/replica/resources/document";
+import {
+  createStaticNotebookBackend,
+  emptySnapshot,
+  type StaticNotebookSnapshot,
+  startWithSnapshot,
+} from "../static";
+import { fixtureEvents } from "../../notebook/replica/__tests__/fixtures";
+
+function snapshot(): StaticNotebookSnapshot {
+  const events = fixtureEvents("notebook-events.sse");
+  const notebook = events.find((event) => event.event === "notebook")?.data
+    .notebook;
+  const document = events.find((event) => event.event === "document")?.data
+    .document;
+  assert(notebook && document);
+  return {
+    notebook,
+    document,
+    cells: events.flatMap((event) =>
+      event.event === "cell" ? [event.data.cell] : [],
+    ),
+  };
+}
+
+const streams: NotebookEvents[] = [];
+afterEach(async () => {
+  const closed = streams.splice(0);
+  closed.forEach((stream) => stream.close());
+  await Promise.all(closed.map((stream) => stream.closed));
+  vi.useRealTimers();
+});
+
+/** Collect a stream's events up to and including its next `ready`. */
+function untilReady(stream: NotebookEvents): Promise<Protocol.EventOutput[]> {
+  const events: Protocol.EventOutput[] = [];
+  return new Promise((resolve, reject) => {
+    stream.subscribe((event) => {
+      events.push(Protocol.Event.parse(event));
+      if (event.event === "ready") {
+        resolve(events);
+      }
+    }, reject);
+  });
+}
+
+/** Open a backend's stream and read until it is ready. */
+async function read(
+  backend: ReturnType<typeof createStaticNotebookBackend>,
+  cursor?: string,
+) {
+  const stream = new NotebookEvents(backend.transport, { cursor });
+  streams.push(stream);
+  const events = await untilReady(stream);
+  return { stream, events };
+}
+
+/** A live backend whose stream sends nothing until `release` is called. */
+function silentUntilReleased(backend: NotebookBackend) {
+  let release = () => {};
+  const released = new Promise<void>((resolve) => {
+    release = resolve;
+  });
+  const attempts: Array<string | null> = [];
+  const silent: NotebookBackend = {
+    ...backend,
+    transport: async (attempt, handlers) => {
+      attempts.push(attempt.cursor);
+      await released;
+      return backend.transport(attempt, handlers);
+    },
+  };
+  return { backend: silent, release: () => release(), attempts };
+}
+
+describe("static notebook backend", () => {
+  it("supplies a complete document through the normal event reader and reducer", async () => {
+    const source = snapshot();
+    const { stream, events } = await read(createStaticNotebookBackend(source));
+    const reducer = createNotebookReducer({ document: documentResource });
+    const state = events.reduce(reducer.reduce, reducer.initial);
+    expect(events.map((event) => event.event)).toEqual([
+      "notebook",
+      "cell",
+      "cell",
+      "cell",
+      "document",
+      "ready",
+    ]);
+    expect(state.notebook).toEqual({ ...source.notebook, operations: [] });
+    expect(state.resources.document.current).toEqual({
+      revision: source.document,
+      cells: new Map(source.cells.map((cell) => [cell.id, cell])),
+    });
+    expect(stream.status).toBe("connected");
+  });
+
+  it("preserves withheld source and its revision metadata", async () => {
+    const source = snapshot();
+    const cells = source.cells.map((cell) => ({ ...cell, code: null }));
+    const { events } = await read(
+      createStaticNotebookBackend({ ...source, cells }),
+    );
+    expect(
+      events
+        .filter((event) => event.event === "cell")
+        .map((event) => event.data.cell),
+    ).toEqual(cells);
+    expect(
+      events.find((event) => event.event === "document")?.data.document,
+    ).toEqual(source.document);
+  });
+
+  it("serves an empty document as available", async () => {
+    const source = snapshot();
+    const document = { ...source.document, order: [] };
+    const { events } = await read(
+      createStaticNotebookBackend({ ...source, document, cells: [] }),
+    );
+    const reducer = createNotebookReducer({ document: documentResource });
+    expect(
+      events.reduce(reducer.reduce, reducer.initial).resources.document.current,
+    ).toEqual({
+      revision: document,
+      cells: new Map(),
+    });
+  });
+
+  it("keeps independent readers open without polling or reconnecting", async () => {
+    vi.useFakeTimers();
+    const backend = createStaticNotebookBackend(snapshot());
+    const first = await read(backend);
+    const second = await read(backend);
+    first.stream.close();
+    await first.stream.closed;
+    await vi.advanceTimersByTimeAsync(60_000);
+    expect(second.events).toEqual(first.events);
+    expect(second.stream.status).toBe("connected");
+    expect(vi.getTimerCount()).toBe(0);
+  });
+
+  it("resets a resume cursor before replaying its snapshot", async () => {
+    const { events } = await read(
+      createStaticNotebookBackend(snapshot()),
+      "previous-10",
+    );
+    expect(events[0]).toEqual({ event: "reset", data: {} });
+    expect(events.at(-1)).toEqual({ event: "ready", id: "static", data: {} });
+  });
+
+  it.each(["edit", "save"] as const)(
+    "rejects direct %s commands",
+    async (operation) => {
+      const backend = createStaticNotebookBackend(snapshot());
+      const client = backend.notebook(snapshot().notebook.id);
+      const result =
+        operation === "edit" ? await client.edit([]) : await client.save();
+      expect(result).toEqual({
+        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",
+        },
+      });
+    },
+  );
+
+  it("respects cancellation before opening, during delivery, and for commands", async () => {
+    const backend = createStaticNotebookBackend(snapshot());
+    const before = new AbortController();
+    before.abort();
+    const open = vi.fn();
+    const message = vi.fn();
+    await backend.transport(
+      { cursor: null, signal: before.signal },
+      { open, message },
+    );
+    expect(open).not.toHaveBeenCalled();
+    expect(message).not.toHaveBeenCalled();
+    await expect(
+      backend
+        .notebook(snapshot().notebook.id)
+        .edit([], { signal: before.signal }),
+    ).rejects.toBe(before.signal.reason);
+    const during = new AbortController();
+    message.mockImplementation(() => during.abort());
+    await backend.transport(
+      { cursor: null, signal: during.signal },
+      { open, message },
+    );
+    expect(message).toHaveBeenCalledOnce();
+  });
+});
+
+describe("starting from an embedded snapshot", () => {
+  it("shows the snapshot until the host answers, then resets into its stream", async () => {
+    const embedded = snapshot();
+    const live = silentUntilReleased(
+      createStaticNotebookBackend(emptySnapshot("live")),
+    );
+    const reducer = createNotebookReducer({ document: documentResource });
+
+    // The snapshot is published on its own, with the host still silent.
+    const { stream, events } = await read(
+      startWithSnapshot(embedded, live.backend),
+    );
+    expect(events.map((event) => event.event)).toEqual([
+      "notebook",
+      "cell",
+      "cell",
+      "cell",
+      "document",
+      "ready",
+    ]);
+    const shown = events.reduce(reducer.reduce, reducer.initial);
+    expect(shown.notebook?.id).toBe(embedded.notebook.id);
+    expect(shown.resources.document.current?.cells.size).toBe(3);
+    expect(live.attempts).toEqual([null]);
+
+    // The host's first frame discards it, and the host's snapshot follows.
+    const host = untilReady(stream);
+    live.release();
+    const rest = await host;
+    expect(rest.map((event) => event.event)).toEqual([
+      "reset",
+      "notebook",
+      "document",
+      "ready",
+    ]);
+    const state = rest.reduce(reducer.reduce, shown);
+    expect(state.notebook?.id).toBe("live");
+    expect(state.resources.document.current?.cells.size).toBe(0);
+    expect(stream.status).toBe("connected");
+  });
+
+  it("does not open the host when cancelled while showing the snapshot", async () => {
+    const attempts: Array<string | null> = [];
+    const live: NotebookBackend = {
+      ...createStaticNotebookBackend(emptySnapshot("live")),
+      transport: async (attempt) => {
+        attempts.push(attempt.cursor);
+      },
+    };
+    const controller = new AbortController();
+    const message = vi.fn(() => controller.abort());
+
+    await startWithSnapshot(snapshot(), live).transport(
+      { cursor: null, signal: controller.signal },
+      { open: vi.fn(), message },
+    );
+
+    expect(message).toHaveBeenCalledOnce();
+    expect(attempts).toEqual([]);
+  });
+
+  it("starts the host connection fresh when resuming from the snapshot's cursor", async () => {
+    const embedded = snapshot();
+    const live = silentUntilReleased(
+      createStaticNotebookBackend(emptySnapshot("live")),
+    );
+    live.release();
+    const { events } = await read(
+      startWithSnapshot(embedded, live.backend),
+      `${embedded.notebook.id}:snapshot`,
+    );
+    expect(live.attempts).toEqual([null]);
+    expect(events.map((event) => event.event)).toEqual([
+      "reset",
+      "notebook",
+      "document",
+      "ready",
+    ]);
+  });
+
+  it("skips the snapshot when a connection resumes", async () => {
+    const backend = startWithSnapshot(
+      snapshot(),
+      createStaticNotebookBackend(emptySnapshot("live")),
+    );
+    const { events } = await read(backend, "static");
+    expect(events.map((event) => event.event)).toEqual([
+      "reset",
+      "notebook",
+      "document",
+      "ready",
+    ]);
+  });
+
+  it("keeps the live backend's commands", async () => {
+    const backend = startWithSnapshot(
+      snapshot(),
+      createStaticNotebookBackend(emptySnapshot("live")),
+    );
+    const result = await backend.notebook("live").save();
+    expect(result).toMatchObject({ ok: false, status: 501 });
+  });
+});

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