@marimo-team/islands 0.25.3-dev34 → 0.25.3-dev349

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

@@ -0,0 +1,340 @@
+/* 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 { MockExecutions } from "@/__mocks__/executions";
+import type { NotebookBackend } from "../../notebook/replica";
+import { document as documentResource } from "../../notebook/replica/resources/document";
+import { results as resultsResource } from "../../notebook/replica/resources/results";
+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] : [],
+    ),
+    results: [],
+  };
+}
+
+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("delivers embedded results as completed executions after the document", async () => {
+    const source = snapshot();
+    const cell = source.cells[0];
+    const result = MockExecutions.completed(cell.id, {
+      id: "embedded:result",
+      run_id: "embedded:run",
+      runtime_id: "embedded",
+      code_fingerprint: cell.code_fingerprint,
+      outcome: {
+        kind: "value",
+        display: { mimetype: "text/html", data: "<b>1</b>" },
+      },
+      started_at: null,
+      completed_at: null,
+    });
+    const { events } = await read(
+      createStaticNotebookBackend({ ...source, results: [result] }),
+    );
+    const reducer = createNotebookReducer({
+      document: documentResource,
+      results: resultsResource,
+    });
+    const state = events.reduce(reducer.reduce, reducer.initial);
+    expect(events.map((event) => event.event).slice(-3)).toEqual([
+      "document",
+      "execution",
+      "ready",
+    ]);
+    expect(state.resources.results.get(cell.id)).toEqual(result);
+  });
+
+  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: "urn:marimo:problem: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: "urn:marimo:problem: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",
         },
       };