@marimo-team/islands 0.25.3-dev141 → 0.25.3-dev142

package/src/core/websocket/__tests__/backend.test.ts

@@ -0,0 +1,235 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+import { NotebookEvents, StreamStopped } from "@marimo-team/client";
+import { createStore } from "jotai/vanilla";
+import { afterEach, describe, expect, it, onTestFinished, vi } from "vitest";
+import { Deferred } from "@/utils/Deferred";
+import { ConnectionState } from "../../host/connection";
+import {
+  type FakeEdit,
+  type FakeSave,
+  fakeBackend,
+} from "../../host/__tests__/fixtures";
+import { createNotebookReplica } from "../../host/replica";
+import { createEditorBackend } from "../backend";
+
+const wasm = vi.hoisted(() => ({ enabled: false, create: vi.fn() }));
+vi.mock("../../wasm/utils", () => ({ isWasm: () => wasm.enabled }));
+vi.mock("../../wasm/bridge", () => ({
+  PyodideBridge: { INSTANCE: { createNotebookBackend: wasm.create } },
+}));
+
+afterEach(() => {
+  wasm.enabled = false;
+  wasm.create.mockReset();
+  vi.unstubAllGlobals();
+});
+
+const options = {
+  static: false,
+  transportType: "websocket" as const,
+  url: () => "ws://localhost/ws",
+  headers: () => ({}),
+  notebookUrl: (id: string) =>
+    new URL(`http://localhost/api/marimo/v1/notebooks/${id}`),
+  waitUntilReady: async () => undefined,
+  checkHealth: async () => true,
+};
+
+/** Record each socket the transport opens. They never connect. */
+function sockets() {
+  const opened: string[] = [];
+  vi.stubGlobal(
+    "WebSocket",
+    class extends EventTarget {
+      public constructor(url: URL) {
+        super();
+        opened.push(String(url));
+      }
+      public close() {}
+    },
+  );
+  return opened;
+}
+
+describe("createEditorBackend", () => {
+  it.each([
+    ["a static notebook", { ...options, static: true }, "idle"],
+    ["WASM", options, "wasm"],
+    ["SSE", { ...options, transportType: "sse" as const }, "fetch"],
+    ["the default", options, "socket"],
+  ] as const)("follows %s", async (_name, chosen, expected) => {
+    wasm.enabled = expected === "wasm";
+    const worker = vi.fn(() => new Promise<void>(() => undefined));
+    wasm.create.mockResolvedValue(fakeBackend({ transport: worker }));
+    const fetch = vi.fn(() => new Promise<Response>(() => undefined));
+    const opened = sockets();
+    vi.stubGlobal("fetch", fetch);
+    const events = new NotebookEvents(createEditorBackend(chosen).transport);
+    await vi.waitFor(() =>
+      expect(
+        {
+          idle: events.status === "connected" && !wasm.create.mock.calls.length,
+          wasm: worker.mock.calls.length === 1,
+          fetch: fetch.mock.calls.length === 1,
+          socket: opened.length === 1,
+        }[expected],
+      ).toBe(true),
+    );
+    events.close();
+  });
+
+  it("edits a reported notebook on its host with the session's headers", async () => {
+    const fetch = vi.fn(
+      async (_input: RequestInfo | URL, _init?: RequestInit) =>
+        new Response(
+          JSON.stringify({ type: "conflict", title: "Conflict", status: 409 }),
+          {
+            status: 409,
+            headers: { "Content-Type": "application/problem+json" },
+          },
+        ),
+    );
+    vi.stubGlobal("fetch", fetch);
+    const backend = createEditorBackend({
+      ...options,
+      headers: () => ({
+        Authorization: "Bearer editor-token",
+        "Marimo-Server-Token": "server-token",
+      }),
+    });
+    const result = await backend
+      .notebook("n1")
+      .edit([], { idempotencyKey: "edit-1" });
+    expect(result).toMatchObject({ ok: false, status: 409 });
+    const [[input, init]] = fetch.mock.calls;
+    expect(String(input)).toBe(
+      "http://localhost/api/marimo/v1/notebooks/n1/edits",
+    );
+    const headers = new Headers(init?.headers);
+    expect(init?.method).toBe("POST");
+    expect(headers.get("Authorization")).toBe("Bearer editor-token");
+    expect(headers.get("Marimo-Server-Token")).toBe("server-token");
+    expect(headers.get("Idempotency-Key")).toBe("edit-1");
+  });
+
+  it("does not edit a static notebook", async () => {
+    const fetch = vi.fn();
+    vi.stubGlobal("fetch", fetch);
+    const result = await createEditorBackend({ ...options, static: true })
+      .notebook("n1")
+      .edit([], { idempotencyKey: "edit-1" });
+    expect(result).toMatchObject({ ok: false, status: 501 });
+    expect(fetch).not.toHaveBeenCalled();
+  });
+
+  it("routes WASM commands through the worker once its stream starts", async () => {
+    wasm.enabled = true;
+    const edit = vi.fn<FakeEdit>();
+    const save = vi.fn<FakeSave>();
+    const transport = vi.fn(async () => undefined);
+    wasm.create.mockResolvedValue(fakeBackend({ transport, edit, save }));
+    const backend = createEditorBackend(options);
+    expect(wasm.create).not.toHaveBeenCalled();
+    expect(() => backend.notebook("n1")).toThrow("has not started");
+
+    const attempt = { cursor: null, signal: new AbortController().signal };
+    const handlers = { open: vi.fn(), message: vi.fn() };
+    await backend.transport(attempt, handlers);
+    await backend.transport(attempt, handlers);
+    expect(wasm.create).toHaveBeenCalledOnce();
+    expect(transport).toHaveBeenCalledWith(attempt, handlers);
+    await backend.notebook("n1").edit([], { idempotencyKey: "edit-1" });
+    expect(edit).toHaveBeenCalledWith("n1", [], { idempotencyKey: "edit-1" });
+    await backend.notebook("n2").save({ idempotencyKey: "save-1" });
+    expect(save).toHaveBeenCalledWith("n2", { idempotencyKey: "save-1" });
+  });
+
+  it("stops following WASM when the worker does not start", async () => {
+    wasm.enabled = true;
+    wasm.create.mockRejectedValue(new Error("Initialization failed"));
+    await expect(
+      createEditorBackend(options).transport(
+        { cursor: null, signal: new AbortController().signal },
+        { open: vi.fn(), message: vi.fn() },
+      ),
+    ).rejects.toBeInstanceOf(StreamStopped);
+  });
+});
+
+describe("the session backend's startup", () => {
+  it.each([true, false])(
+    "rechecks health after failed startup (healthy: %s)",
+    async (healthy) => {
+      const opened = sockets();
+      const waitUntilReady = vi
+        .fn()
+        .mockRejectedValue(new Error("Startup gave up"));
+      const checkHealth = vi.fn().mockResolvedValue(healthy);
+      const store = createStore();
+      const client = createNotebookReplica(store);
+      onTestFinished(
+        client.connect(
+          createEditorBackend({ ...options, waitUntilReady, checkHealth }),
+          { onNotification: vi.fn() },
+        ),
+      );
+      await vi.waitFor(() =>
+        expect(store.get(client.connectionAtom).state).toBe(
+          ConnectionState.CLOSED,
+        ),
+      );
+
+      await client.reconnect();
+      await vi.waitFor(() => expect(opened).toHaveLength(healthy ? 1 : 0));
+      expect({
+        healthChecks: checkHealth.mock.calls.length,
+        startupWaits: waitUntilReady.mock.calls.length,
+        state: store.get(client.connectionAtom).state,
+      }).toEqual({
+        healthChecks: 1,
+        startupWaits: 1,
+        state: healthy ? ConnectionState.CONNECTING : ConnectionState.CLOSED,
+      });
+    },
+  );
+
+  it("shares one startup wait across connections", async () => {
+    const opened = sockets();
+    const ready = new Deferred<void>();
+    const waitUntilReady = vi.fn(() => ready.promise);
+    const backend = createEditorBackend({ ...options, waitUntilReady });
+    const client = createNotebookReplica(createStore());
+    // Strict Mode connects, disconnects, and connects again.
+    client.connect(backend, { onNotification: vi.fn() })();
+    onTestFinished(client.connect(backend, { onNotification: vi.fn() }));
+    await Promise.resolve();
+    expect(opened).toEqual([]);
+
+    ready.resolve();
+    await vi.waitFor(() => expect(opened).toHaveLength(1));
+    expect(waitUntilReady).toHaveBeenCalledOnce();
+  });
+
+  it("opens nothing for a connection that ended during startup", async () => {
+    const opened = sockets();
+    const ready = new Deferred<void>();
+    const store = createStore();
+    const client = createNotebookReplica(store);
+    const disconnect = client.connect(
+      createEditorBackend({ ...options, waitUntilReady: () => ready.promise }),
+      { onNotification: vi.fn() },
+    );
+    expect(store.get(client.connectionAtom)).toEqual({
+      state: ConnectionState.CONNECTING,
+    });
+    disconnect();
+    ready.resolve();
+    await ready.promise;
+    await Promise.resolve();
+    expect(opened).toEqual([]);
+    expect(store.get(client.connectionAtom)).toEqual({
+      state: ConnectionState.NOT_STARTED,
+    });
+  });
+});

package/src/core/websocket/__tests__/{useMarimoKernelConnection.test.ts → connection-state.test.ts}

@@ -3,7 +3,7 @@
 import { afterEach, describe, expect, it, vi } from "vitest";
 import { Logger } from "@/utils/Logger";
 import { WebSocketClosedReason, ConnectionState } from "../types";
-import { classifyCloseEvent } from "../useMarimoKernelConnection";
+import { classifyCloseEvent } from "../connection-state";
 
 function classify(reason: string | undefined) {
   return classifyCloseEvent({ reason });
@@ -13,8 +13,10 @@ describe("classifyCloseEvent", () => {
   describe("transient closes (default branch)", () => {
     it("retries on empty/undefined reason", () => {
       const decision = classify(undefined);
-      expect(decision.kind).toBe("retry");
-      expect(decision.status).toEqual({ state: ConnectionState.CONNECTING });
+      expect(decision).toEqual({
+        kind: "retry",
+        status: { state: ConnectionState.CONNECTING },
+      });
     });
 
     it("treats unknown reason strings as transient and logs a warning", () => {
@@ -40,41 +42,25 @@ describe("classifyCloseEvent", () => {
       // terminal or a rejected client would reconnect forever.
       "MARIMO_UNAUTHORIZED",
       "MARIMO_KIOSK_NOT_ALLOWED",
-    ])("%s → terminal with KERNEL_DISCONNECTED, closes transport", (reason) => {
+    ])("%s → terminal with KERNEL_DISCONNECTED", (reason) => {
       const decision = classify(reason);
-      expect(decision.kind).toBe("terminal");
-      expect(decision.status).toMatchObject({
-        state: ConnectionState.CLOSED,
-        code: WebSocketClosedReason.KERNEL_DISCONNECTED,
-      });
       expect(decision).toMatchObject({
         kind: "terminal",
-        closeTransport: true,
+        status: {
+          state: ConnectionState.CLOSED,
+          code: WebSocketClosedReason.KERNEL_DISCONNECTED,
+        },
       });
     });
 
-    it("MARIMO_KERNEL_STARTUP_ERROR → terminal + closeTransport", () => {
+    it("MARIMO_KERNEL_STARTUP_ERROR → terminal", () => {
       const decision = classify("MARIMO_KERNEL_STARTUP_ERROR");
-      expect(decision.kind).toBe("terminal");
-      expect(decision.status).toMatchObject({
-        state: ConnectionState.CLOSED,
-        code: WebSocketClosedReason.KERNEL_STARTUP_ERROR,
-      });
       expect(decision).toMatchObject({
         kind: "terminal",
-        closeTransport: true,
-      });
-    });
-  });
-
-  describe("transport exhaustion", () => {
-    it("MARIMO_TRANSPORT_EXHAUSTED → gave-up with KERNEL_DISCONNECTED", () => {
-      const decision = classify("MARIMO_TRANSPORT_EXHAUSTED");
-      expect(decision.kind).toBe("gave-up");
-      expect(decision.status).toEqual({
-        state: ConnectionState.CLOSED,
-        code: WebSocketClosedReason.KERNEL_DISCONNECTED,
-        reason: "kernel not found",
+        status: {
+          state: ConnectionState.CLOSED,
+          code: WebSocketClosedReason.KERNEL_STARTUP_ERROR,
+        },
       });
     });
   });

package/src/core/websocket/__tests__/useNotebookReplica.hook.test.tsx

@@ -0,0 +1,141 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+// @vitest-environment jsdom
+
+import { act, renderHook } from "@testing-library/react";
+import { createStore, Provider } from "jotai";
+import type { PropsWithChildren } from "react";
+import { beforeEach, describe, expect, it, vi } from "vitest";
+import type { NotebookTransport } from "@marimo-team/client";
+import { ConnectionState } from "../types";
+import { createEditorBackend } from "../backend";
+import { useNotebookReplica } from "../useNotebookReplica";
+import { fakeBackend, fixtureMessages } from "../../host/__tests__/fixtures";
+import { agentAttachmentsAtom } from "../../attachments/state";
+import { notebookReplicaAtom } from "../state";
+
+vi.mock("../backend", () => ({ createEditorBackend: vi.fn() }));
+
+/** A host-framed native notification. */
+const reconnected = {
+  event: "marimo",
+  data: { marimo: { op: "reconnected" } },
+};
+
+/** A transport whose connections the test drives. */
+function makeTransport() {
+  const opened: {
+    signal: AbortSignal;
+    open(): void;
+    receive(message: unknown): void;
+  }[] = [];
+  const transport: NotebookTransport = ({ signal }, handlers) =>
+    new Promise((resolve) => {
+      opened.push({ signal, open: handlers.open, receive: handlers.message });
+      signal.addEventListener("abort", () => resolve(), { once: true });
+    });
+  return {
+    transport,
+    opened,
+    /** Open the latest connection and deliver a host-framed notification. */
+    notify: () => opened.at(-1)?.receive(reconnected),
+  };
+}
+
+function setup(strict = false) {
+  const store = createStore();
+  const options = {
+    static: false,
+    transportType: "websocket" as const,
+    url: () => "ws://localhost/ws",
+    headers: () => ({}),
+    notebookUrl: (id: string) => new URL(`http://localhost/notebooks/${id}`),
+    waitUntilReady: async () => undefined,
+    checkHealth: async () => true,
+    onNotification: vi.fn(),
+  };
+  const wrapper = ({ children }: PropsWithChildren) => (
+    <Provider store={store}>{children}</Provider>
+  );
+  const hook = renderHook((props) => useNotebookReplica(props), {
+    initialProps: options,
+    reactStrictMode: strict,
+    wrapper,
+  });
+  return { ...hook, options, store };
+}
+
+beforeEach(() => {
+  vi.mocked(createEditorBackend).mockReset();
+});
+
+describe("useNotebookReplica", () => {
+  it("exposes host presence to the UI without waiting for kernel-ready", async () => {
+    const fake = makeTransport();
+    vi.mocked(createEditorBackend).mockReturnValue(
+      fakeBackend({ transport: fake.transport }),
+    );
+    const { result, store, unmount } = setup();
+    await act(async () => Promise.resolve());
+    expect(store.get(notebookReplicaAtom)).toBe(result.current);
+    const [connection] = fake.opened;
+    act(() => {
+      connection.open();
+      fixtureMessages("notebook-events.sse").forEach(connection.receive);
+    });
+    expect(store.get(agentAttachmentsAtom)).toEqual([
+      {
+        id: "a2",
+        notebook_id: "n1",
+        kind: "agent",
+        name: "Notebook CLI",
+        since: "2026-10-01T09:20:03Z",
+      },
+    ]);
+    unmount();
+    expect(store.get(notebookReplicaAtom)).toBeNull();
+    expect(store.get(agentAttachmentsAtom)).toEqual([]);
+  });
+
+  it("uses current callbacks and URL providers without reconnecting on a rerender", async () => {
+    const fake = makeTransport();
+    vi.mocked(createEditorBackend).mockReturnValue(
+      fakeBackend({ transport: fake.transport }),
+    );
+    const { rerender, options, unmount } = setup();
+    await act(async () => Promise.resolve());
+    const onNotification = vi.fn();
+    rerender({ ...options, onNotification, url: () => "ws://redirected/ws" });
+    act(() => fake.notify());
+    expect(options.onNotification).not.toHaveBeenCalled();
+    expect(onNotification).toHaveBeenCalledExactlyOnceWith({
+      data: { op: "reconnected" },
+    });
+    const [{ url }] = vi.mocked(createEditorBackend).mock.calls[0];
+    expect(url()).toBe("ws://redirected/ws");
+    expect(fake.opened).toHaveLength(1);
+    unmount();
+  });
+
+  it("follows only the active connection under Strict Mode", async () => {
+    const fake = makeTransport();
+    vi.mocked(createEditorBackend).mockReturnValue(
+      fakeBackend({ transport: fake.transport }),
+    );
+    const { result, store, options, unmount } = setup(true);
+    await act(async () => Promise.resolve());
+    const live = fake.opened.filter((attempt) => !attempt.signal.aborted);
+    expect(live).toHaveLength(1);
+    const [active] = live;
+    act(() => active.open());
+    expect(store.get(result.current.connectionAtom)).toEqual({
+      state: ConnectionState.OPEN,
+    });
+    act(() => fake.opened.forEach((attempt) => attempt.receive(reconnected)));
+    expect(options.onNotification).toHaveBeenCalledOnce();
+    unmount();
+    expect(active.signal.aborted).toBe(true);
+    const attempts = fake.opened.length;
+    await result.current.reconnect();
+    expect(fake.opened).toHaveLength(attempts);
+  });
+});