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

package/src/core/saving/state.ts

@@ -1,52 +1,51 @@
 /* Copyright 2026 Marimo. All rights reserved. */
 
+/** Whether the notebook has changes the host has not written to its store. */
+
 import { dequal as isEqual } from "dequal";
 import { atom } from "jotai";
-import { arrayShallowEquals } from "@/utils/arrays";
-import { type NotebookState, notebookAtom } from "../cells/cells";
 import { type LayoutState, layoutStateAtom } from "../layout/state";
-import type { CellConfig } from "../network/types";
-
-export interface LastSavedNotebook {
-  codes: string[];
-  configs: CellConfig[];
-  names: string[];
-  layout: LayoutState;
-}
+import {
+  activeDocumentModelAtom,
+  confirmedDocumentAtom,
+} from "../session/active-notebook";
 
-export const lastSavedNotebookAtom = atom<LastSavedNotebook | undefined>(
-  undefined,
+/** The document version, which counts confirmed changes. Null without a document. */
+export const documentVersionAtom = atom(
+  (get) => get(confirmedDocumentAtom)?.revision.version ?? null,
 );
 
-export const needsSaveAtom = atom((get) => {
-  const lastSavedNotebook = get(lastSavedNotebookAtom);
-  const state = get(notebookAtom);
-  const layout = get(layoutStateAtom);
-  return notebookNeedsSave({ state, layout, lastSavedNotebook });
-});
+/**
+ * The layout as the file last had it. The host document does not carry
+ * the layout file, so this client compares it itself until it does.
+ */
+export const savedLayoutAtom = atom<LayoutState | undefined>(undefined);
 
-function notebookNeedsSave({
-  state,
-  layout,
-  lastSavedNotebook,
-}: {
-  state: NotebookState;
-  layout: LayoutState;
-  lastSavedNotebook: LastSavedNotebook | undefined;
-}) {
-  if (!lastSavedNotebook) {
+/**
+ * Whether saving would change the file: the host's document differs from
+ * what it last wrote, this client has edits the host has not confirmed,
+ * or the layout changed. A document never written reports as unsaved.
+ */
+export const needsSaveAtom = atom((get) => {
+  const document = get(confirmedDocumentAtom);
+  if (document === null) {
     return false;
   }
-  const { cellIds, cellData } = state;
-  const data = cellIds.inOrderIds.map((cellId) => cellData[cellId]);
-  const codes = data.map((d) => d.code);
-  const configs = data.map((d) => d.config);
-  const names = data.map((d) => d.name);
+  if (document.revision.etag !== document.revision.saved_etag) {
+    return true;
+  }
+  const model = get(activeDocumentModelAtom);
+  if (model !== null) {
+    const { status } = get(model.syncStatusAtom);
+    if (status !== "idle" && status !== "unavailable") {
+      return true;
+    }
+  }
+  const saved = get(savedLayoutAtom);
+  const layout = get(layoutStateAtom);
   return (
-    !arrayShallowEquals(codes, lastSavedNotebook.codes) ||
-    !arrayShallowEquals(names, lastSavedNotebook.names) ||
-    !isEqual(configs, lastSavedNotebook.configs) ||
-    !isEqual(layout.selectedLayout, lastSavedNotebook.layout.selectedLayout) ||
-    !isEqual(layout.layoutData, lastSavedNotebook.layout.layoutData)
+    saved !== undefined &&
+    (layout.selectedLayout !== saved.selectedLayout ||
+      !isEqual(layout.layoutData, saved.layoutData))
   );
-}
+});

package/src/core/saving/useAutoSave.ts

@@ -1,68 +1,32 @@
 /* Copyright 2026 Marimo. All rights reserved. */
-import { useEffect, useRef } from "react";
+import { useEffect } from "react";
 import type { UserConfig } from "../config/config-schema";
-import type { CellConfig } from "../network/types";
-import { type ConnectionStatus, ConnectionState } from "../websocket/types";
+import { type ConnectionStatus, ConnectionState } from "../session/types";
 
+/** Save after a delay while the notebook needs it, restarting on every confirmed change. */
 export function useAutoSave(opts: {
-  codes: string[];
-  cellConfigs: CellConfig[];
-  cellNames: string[];
+  /** The document version, which counts confirmed changes. */
+  version: number | null;
   config: UserConfig["save"];
   connStatus: ConnectionStatus;
   needsSave: boolean;
-  kioskMode: boolean;
+  canEdit: boolean;
   onSave: () => void;
 }) {
-  const {
-    codes,
-    cellConfigs,
-    config,
-    connStatus,
-    cellNames,
-    needsSave,
-    kioskMode,
-    onSave,
-  } = opts;
-  const autosaveTimeoutId = useRef<NodeJS.Timeout | null>(null);
-
-  const codesString = codes.join(":");
-  const cellConfigsString = cellConfigs
-    .map((config) => JSON.stringify(config))
-    .join(":");
-  const cellNamesString = cellNames.join(":");
+  const { version, config, connStatus, needsSave, canEdit, onSave } = opts;
 
   useEffect(() => {
-    // If kiosk mode is enabled, do not autosave
-    if (kioskMode) {
+    if (
+      !canEdit ||
+      config.autosave !== "after_delay" ||
+      !needsSave ||
+      connStatus.state !== ConnectionState.OPEN
+    ) {
       return;
     }
 
-    if (config.autosave === "after_delay") {
-      if (autosaveTimeoutId.current !== null) {
-        clearTimeout(autosaveTimeoutId.current);
-      }
-
-      if (needsSave && connStatus.state === ConnectionState.OPEN) {
-        autosaveTimeoutId.current = setTimeout(onSave, config.autosave_delay);
-      }
-    }
-
-    return () => {
-      if (autosaveTimeoutId.current !== null) {
-        clearTimeout(autosaveTimeoutId.current);
-      }
-    };
-    // codes, cellConfigs, cellNames needed in deps array to prevent race condition
-    // with needsSave when user changes state rapidly
-  }, [
-    codesString,
-    cellConfigsString,
-    cellNamesString,
-    config,
-    connStatus.state,
-    onSave,
-    kioskMode,
-    needsSave,
-  ]);
+    const timeout = setTimeout(onSave, config.autosave_delay);
+    return () => clearTimeout(timeout);
+    // `version` restarts the delay on every confirmed change.
+  }, [version, config, connStatus.state, onSave, canEdit, needsSave]);
 }

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

@@ -0,0 +1,236 @@
+/* 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 "../../notebook/replica/connection";
+import {
+  type FakeEdit,
+  type FakeSave,
+  fakeBackend,
+} from "../../notebook/replica/__tests__/fixtures";
+import { createNotebookReplica } from "../../notebook/replica/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,
+  snapshot: null,
+  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/session/__tests__/useNotebookReplica.hook.test.tsx

@@ -0,0 +1,184 @@
+/* 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 "../../notebook/replica/__tests__/fixtures";
+import { useNotebookAttachments } from "@/core/notebook";
+import { activeNotebookReplicaAtom } from "../active-notebook";
+
+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,
+    snapshot: null,
+    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();
+    const attachments = renderHook(
+      () => ({
+        all: useNotebookAttachments(),
+        agents: useNotebookAttachments("agent"),
+      }),
+      {
+        wrapper: ({ children }: PropsWithChildren) => (
+          <Provider store={store}>{children}</Provider>
+        ),
+      },
+    );
+    await act(async () => Promise.resolve());
+    expect(store.get(activeNotebookReplicaAtom)).toBe(result.current);
+    const [connection] = fake.opened;
+    act(() => {
+      connection.open();
+      fixtureMessages("notebook-events.sse").forEach(connection.receive);
+    });
+    expect(attachments.result.current.agents).toEqual([
+      {
+        id: "a2",
+        notebook_id: "n1",
+        kind: "agent",
+        name: "Notebook CLI",
+        since: "2026-10-01T09:20:03Z",
+      },
+    ]);
+    expect(attachments.result.current.all).toEqual(
+      store.get(result.current.attachmentsAtom),
+    );
+    expect(attachments.result.current.all.length).toBeGreaterThan(
+      attachments.result.current.agents.length,
+    );
+    const agents = attachments.result.current.agents;
+    const browser = attachments.result.current.all.find(
+      (attachment) => attachment.kind === "browser",
+    )!;
+    act(() =>
+      connection.receive({
+        event: "attachment",
+        data: { attachment: { ...browser, name: "Firefox" } },
+      }),
+    );
+    expect(attachments.result.current.agents).toBe(agents);
+
+    const renamedAgent = { ...agents[0], name: "Renamed agent" };
+    act(() =>
+      connection.receive({
+        event: "attachment",
+        data: { attachment: renamedAgent },
+      }),
+    );
+    expect(attachments.result.current.agents).toEqual([renamedAgent]);
+
+    unmount();
+    expect(store.get(activeNotebookReplicaAtom)).toBeNull();
+    expect(attachments.result.current).toEqual({ all: [], agents: [] });
+    attachments.unmount();
+  });
+
+  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);
+  });
+});