@marimo-team/islands 0.25.3-dev2 → 0.25.3-dev200

package/src/core/saving/save-component.tsx

@@ -1,6 +1,6 @@
 /* Copyright 2026 Marimo. All rights reserved. */
 
-import { useAtom, useAtomValue, useSetAtom, useStore } from "jotai";
+import { useAtom, useAtomValue, useStore } from "jotai";
 import { SaveIcon } from "lucide-react";
 import { useState } from "react";
 import { FilenameInput } from "@/components/editor/header/filename-input";
@@ -20,27 +20,24 @@ import {
 import { Label } from "../../components/ui/label";
 import { useEvent } from "../../hooks/useEvent";
 import { Logger } from "../../utils/Logger";
-import { getCellConfigs, getNotebook, useNotebook } from "../cells/cells";
-import { withDocumentSave } from "../cells/document-changes";
+import { getCellConfigs, getNotebook } from "../cells/cells";
+import { flushDocument, withDocumentSave } from "../cells/document";
 import { notebookCells } from "../cells/utils";
 import { formatAll } from "../codemirror/format";
 import { autoSaveConfigAtom } from "../config/config";
 import { useAutoExport } from "../export/hooks";
 import { getSerializedLayout } from "../layout/layout";
 import { layoutStateAtom } from "../layout/state";
-import { kioskModeAtom } from "../mode";
 import { connectionAtom } from "../network/connection";
+import { canEditAtom } from "@/core/notebook";
 import { useRequestClient } from "../network/requests";
-import { ConnectionState } from "../websocket/types";
+import { ConnectionState } from "../session/types";
+import { isWasm } from "../wasm/utils";
 import { filenameAtom } from "./file-state";
 import { useFilename, useUpdateFilename } from "./filename";
-import { lastSavedNotebookAtom, needsSaveAtom } from "./state";
+import { documentVersionAtom, needsSaveAtom, savedLayoutAtom } from "./state";
 import { useAutoSave } from "./useAutoSave";
 
-interface SaveNotebookProps {
-  kioskMode: boolean;
-}
-
 let notebookSaveQueue: Promise<void> = Promise.resolve();
 
 export function enqueueNotebookSave<T>(save: () => Promise<T>): Promise<T> {
@@ -52,12 +49,12 @@ export function enqueueNotebookSave<T>(save: () => Promise<T>): Promise<T> {
   return result;
 }
 
-export const SaveComponent = ({ kioskMode }: SaveNotebookProps) => {
+export const SaveComponent = () => {
   const filename = useFilename();
   const needsSave = useAtomValue(needsSaveAtom);
   const closed = useAtomValue(connectionAtom).state === ConnectionState.CLOSED;
   const { saveOrNameNotebook, saveIfNotebookIsPersistent } = useSaveNotebook();
-  useAutoSaveNotebook({ onSave: saveIfNotebookIsPersistent, kioskMode });
+  useAutoSaveNotebook({ onSave: saveIfNotebookIsPersistent });
 
   useAutoExport();
 
@@ -104,7 +101,6 @@ export const SaveComponent = ({ kioskMode }: SaveNotebookProps) => {
 export function useSaveNotebook() {
   const { sendSave } = useRequestClient();
   const { openModal, closeModal, openAlert } = useImperativeModal();
-  const setLastSavedNotebook = useSetAtom(lastSavedNotebookAtom);
   const updateFilename = useUpdateFilename();
   const store = useStore();
 
@@ -113,11 +109,6 @@ export function useSaveNotebook() {
     enqueueNotebookSave(async () => {
       const connection = store.get(connectionAtom);
       const autoSaveConfig = store.get(autoSaveConfigAtom);
-      const kioskMode = store.get(kioskModeAtom);
-
-      if (kioskMode) {
-        return;
-      }
 
       // Don't save if we are not connected to a kernel
       if (connection.state !== ConnectionState.OPEN) {
@@ -125,6 +116,11 @@ export function useSaveNotebook() {
         return;
       }
 
+      // The host may not let this consumer edit, or may not have said yet.
+      if (!store.get(canEditAtom)) {
+        return;
+      }
+
       Logger.log("saving to ", filename);
 
       if (userInitiated && autoSaveConfig.format_on_save) {
@@ -136,7 +132,7 @@ export function useSaveNotebook() {
         return;
       }
 
-      const savedNotebook = await withDocumentSave(async () => {
+      const save = async () => {
         // Grab the latest notebook state after formatting and after every
         // earlier document transaction has reached the server.
         const notebook = getNotebook();
@@ -156,9 +152,14 @@ export function useSaveNotebook() {
           layout: getSerializedLayout(),
           persist: true,
         });
-        return { names: cellNames, codes, configs, layout };
-      });
-      setLastSavedNotebook(savedNotebook);
+        return layout;
+      };
+      // TODO(MO-7955): Remove this conditional once all frontend saves use
+      // the host protocol. Host saves cannot recover edits from an editor snapshot.
+      const savedLayout = isWasm()
+        ? await flushDocument().then(save)
+        : await withDocumentSave(save);
+      store.set(savedLayoutAtom, savedLayout);
     }),
   );
 
@@ -218,29 +219,20 @@ export function isNamedPersistentFile(
   );
 }
 
-export function useAutoSaveNotebook(opts: {
-  onSave: () => void;
-  kioskMode: boolean;
-}) {
+export function useAutoSaveNotebook(opts: { onSave: () => void }) {
   const autoSaveConfig = useAtomValue(autoSaveConfigAtom);
-  const notebook = useNotebook();
+  const canEdit = useAtomValue(canEditAtom);
   const [connection] = useAtom(connectionAtom);
   const needsSave = useAtomValue(needsSaveAtom);
-
-  const cells = notebookCells(notebook);
-  const codes = cells.map((cell) => cell.code);
-  const cellNames = cells.map((cell) => cell.name);
-  const configs = getCellConfigs(notebook);
+  const version = useAtomValue(documentVersionAtom);
 
   useAutoSave({
     onSave: opts.onSave,
     needsSave: needsSave,
-    codes: codes,
-    cellConfigs: configs,
-    cellNames: cellNames,
+    version,
     connStatus: connection,
     config: autoSaveConfig,
-    kioskMode: opts.kioskMode,
+    canEdit,
   });
 }
 

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