@marimo-team/islands 0.25.3-dev28 → 0.25.3-dev284

package/src/core/runtime/__tests__/runtime.test.ts

@@ -2,11 +2,7 @@
 
 import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
 import { Mocks } from "@/__mocks__/common";
-import {
-  generateSessionId,
-  isSessionId,
-  type SessionId,
-} from "@/core/kernel/session";
+import type { SessionId } from "@/core/kernel/session";
 import { Logger } from "@/utils/Logger";
 import { RuntimeManager } from "../runtime";
 import type { RuntimeConfig } from "../types";
@@ -56,38 +52,6 @@ describe("RuntimeManager", () => {
   });
 
   describe("getWsURL", () => {
-    it("uses host notebook events with the existing session and auth parameters", () => {
-      const original = window.location.href;
-      window.history.replaceState({}, "", "/?types=run");
-      try {
-        const runtime = new RuntimeManager({
-          ...mockConfig,
-          url: "https://example.com/nested?file=notebook.py",
-          notebookId: "nb_123",
-        });
-        const sessionId = "s_123456";
-        if (!isSessionId(sessionId)) {
-          throw new Error("Invalid test session ID");
-        }
-        const url = runtime.getWsURL(sessionId);
-        expect({
-          endpoint: `${url.origin}${url.pathname}`,
-          params: Object.fromEntries(url.searchParams),
-        }).toEqual({
-          endpoint:
-            "wss://example.com/nested/api/marimo/v1/notebooks/nb_123/events/ws",
-          params: {
-            file: "notebook.py",
-            types: "notebook,document,cell,attachment,marimo",
-            session_id: sessionId,
-            access_token: "test-token",
-          },
-        });
-      } finally {
-        window.history.replaceState({}, "", original);
-      }
-    });
-
     it("should return WebSocket URL with session ID for https", () => {
       const runtime = new RuntimeManager(mockConfig);
       const sessionId = "1234" as SessionId;
@@ -179,13 +143,9 @@ describe("RuntimeManager", () => {
       );
 
       const wsUrl = runtime.getWsURL("s_123" as SessionId);
-      const wsSyncUrl = runtime.getWsSyncURL("s_123" as SessionId);
       const terminalUrl = runtime.getTerminalWsURL();
 
       expect(wsUrl.searchParams.get("access_token")).toBe("my-secret-token");
-      expect(wsSyncUrl.searchParams.get("access_token")).toBe(
-        "my-secret-token",
-      );
       expect(terminalUrl.searchParams.get("access_token")).toBe(
         "my-secret-token",
       );
@@ -256,22 +216,24 @@ describe("RuntimeManager", () => {
     });
   });
 
-  describe("getSseURL", () => {
-    it.each([undefined, "nb_123"])(
-      "uses the legacy session stream when notebookId is %s",
-      (notebookId) => {
-        const runtime = new RuntimeManager({
-          ...mockConfig,
-          url: "https://example.com/nested?file=notebook.py",
-          notebookId,
-        });
-        const sessionId = generateSessionId();
-        expect(runtime.getSseURL(sessionId).toString()).toBe(
-          `https://example.com/nested/sse?file=notebook.py&session_id=${sessionId}`,
-        );
-      },
-    );
+  describe("getNotebookURL", () => {
+    it("addresses the notebook below the host API without the auth token", () => {
+      const runtime = new RuntimeManager(
+        {
+          url: "https://sandbox.example.com/base",
+          lazy: true,
+          authToken: "my-secret-token",
+        },
+        true,
+      );
+      const url = runtime.getNotebookURL("n/1");
 
+      expect(url.pathname).toBe("/base/api/marimo/v1/notebooks/n%2F1");
+      expect(url.toString()).not.toContain("my-secret-token");
+    });
+  });
+
+  describe("getSseURL", () => {
     it("should never put the auth token in the URL, even cross-origin", () => {
       // SSE requests can send headers, so auth travels in the
       // Authorization header; a token in the URL would leak into
@@ -300,25 +262,14 @@ describe("RuntimeManager", () => {
         const url = runtime.getSseURL("s_123" as SessionId);
 
         expect(url.searchParams.get("access_token")).toBeNull();
-        expect(url.searchParams.get("foo")).toBe("bar");
+        expect(url.searchParams.get("query_params")).toBe("foo=bar");
+        expect(url.toString()).not.toContain("page-token");
       } finally {
         window.history.pushState({}, "", original);
       }
     });
   });
 
-  describe("getWsSyncURL", () => {
-    it("should return WebSocket Sync URL", () => {
-      const runtime = new RuntimeManager(mockConfig);
-      const sessionId = "1234" as SessionId;
-      const url = runtime.getWsSyncURL(sessionId);
-
-      expect(url.protocol).toBe("wss:");
-      expect(url.pathname).toBe("/ws_sync");
-      expect(url.searchParams.get("session_id")).toBe(sessionId);
-    });
-  });
-
   describe("getTerminalWsURL", () => {
     it("should include dimensions in the terminal handshake", () => {
       const runtime = new RuntimeManager(mockConfig);
@@ -720,9 +671,16 @@ describe("RuntimeManager", () => {
       expect(httpUrl.searchParams.get("base_param")).toBe("existing");
 
       // Should preserve all window-level query params (including custom ones)
-      expect(wsUrl.searchParams.get("custom_param")).toBe("value123");
-      expect(wsUrl.searchParams.get("user_token")).toBe("abc");
-      expect(wsUrl.searchParams.get("theme")).toBe("dark");
+      expect(
+        Object.fromEntries(
+          new URLSearchParams(wsUrl.searchParams.get("query_params") ?? ""),
+        ),
+      ).toEqual({
+        base_param: "existing",
+        custom_param: "value123",
+        user_token: "abc",
+        theme: "dark",
+      });
       expect(httpUrl.searchParams.get("custom_param")).toBe("value123");
       expect(httpUrl.searchParams.get("user_token")).toBe("abc");
       expect(httpUrl.searchParams.get("theme")).toBe("dark");
@@ -781,7 +739,7 @@ describe("RuntimeManager", () => {
       );
     });
 
-    it("should throw when creating WebSocket URLs from blob URLs", () => {
+    it("warns when creating WebSocket URLs from blob URLs", () => {
       const runtime = new RuntimeManager({
         url: "blob:https://example.com/12345678-1234-1234-1234-123456789abc",
         lazy: true,
@@ -789,7 +747,7 @@ describe("RuntimeManager", () => {
       const sessionId = "test" as SessionId;
 
       expect(runtime.getWsURL(sessionId).toString()).toMatchInlineSnapshot(
-        `"blob:https://example.com/12345678-1234-1234-1234-123456789abc?session_id=test"`,
+        `"blob:https://example.com/12345678-1234-1234-1234-123456789abc?query_params=&session_id=test"`,
       );
       expect(Logger.warn).toHaveBeenCalledOnce();
     });

package/src/core/runtime/adapter.ts

@@ -1,14 +1,14 @@
 /* Copyright 2026 Marimo. All rights reserved. */
 
 import { type Atom, atom } from "jotai";
-import { kernelStateAtom } from "../kernel/state";
+import { runtimeStatusAtom } from "../executions";
 import { islandsInitializedAtom } from "../islands/state";
 import { isIslands } from "../islands/utils";
 import { connectionAtom } from "../network/connection";
 import { isStaticNotebook } from "../static/static-state";
 import { wasmInitStateAtom } from "../wasm/state";
 import { isWasm } from "../wasm/utils";
-import { ConnectionState } from "../websocket/types";
+import { ConnectionState } from "../session/types";
 
 export type RuntimeKind = "wasm" | "remote" | "static" | "islands";
 
@@ -163,8 +163,8 @@ function selectAdapter(): RuntimeAdapter {
 export const runtimeAdapterAtom = atom<RuntimeAdapter>(selectAdapter());
 
 /**
- * Remote runtimes require both the WS to be open *and* the kernel to have
- * ack'd instantiation. WASM/islands have no separate kernel handshake.
+ * Remote runtimes require both the WS to be open *and* the host to report
+ * the runtime running. WASM/islands have no separate kernel handshake.
  * Static notebooks never run cells.
  */
 export const canRunCellsAtom = atom<boolean>((get) => {
@@ -176,7 +176,7 @@ export const canRunCellsAtom = atom<boolean>((get) => {
     return false;
   }
   if (adapter.kind === "remote") {
-    return get(kernelStateAtom).isInstantiated;
+    return get(runtimeStatusAtom) === "running";
   }
   return true;
 });

package/src/core/runtime/config.ts

@@ -4,8 +4,8 @@ import useEvent from "react-use-event-hook";
 import { Logger } from "@/utils/Logger";
 import { connectionAtom } from "../network/connection";
 import { store } from "../state/jotai";
-import { isAppNotStarted } from "../websocket/connection-utils";
-import { ConnectionState } from "../websocket/types";
+import { isAppNotStarted } from "../session/connection-utils";
+import { ConnectionState } from "../session/types";
 import { RuntimeManager } from "./runtime";
 import type { RuntimeConfig } from "./types";
 

package/src/core/runtime/runtime.ts

@@ -1,15 +1,17 @@
 /* Copyright 2026 Marimo. All rights reserved. */
 
+import type { NotebookId } from "@marimo-team/client/protocol";
 import { Deferred } from "@/utils/Deferred";
 import { Logger } from "@/utils/Logger";
 import { KnownQueryParams } from "../constants";
-import { NOTEBOOK_EVENT_TYPES } from "../host/events";
 import { isIslands } from "../islands/utils";
 import { getSessionId, type SessionId } from "../kernel/session";
 import { isStaticNotebook } from "../static/static-state";
 import { isWasm } from "../wasm/utils";
 import type { RuntimeConfig } from "./types";
 
+const knownQueryParams = new Set(Object.values(KnownQueryParams));
+
 export class RuntimeManager {
   private initialHealthyCheck = new Deferred<void>();
   private config: RuntimeConfig;
@@ -72,10 +74,7 @@ export class RuntimeManager {
     }
 
     for (const [key, value] of currentParams.entries()) {
-      if (
-        restrictToKnownQueryParams &&
-        !Object.values(KnownQueryParams).includes(key)
-      ) {
+      if (restrictToKnownQueryParams && !knownQueryParams.has(key)) {
         continue;
       }
       baseUrl.searchParams.set(key, value);
@@ -130,43 +129,32 @@ export class RuntimeManager {
     return asWsUrl(this.appendCrossOriginAuth(url).toString());
   }
 
-  /**
-   * Search params for a session connection: the base URL's params, merged
-   * with the current page's, plus the session id.
-   */
-  private getSessionSearchParams(sessionId: SessionId): URLSearchParams {
-    const baseUrl = new URL(this.config.url);
-    const searchParams = new URLSearchParams(baseUrl.search);
-
-    // Merge in current page's query parameters
-    const currentParams = new URLSearchParams(window.location.search);
-    currentParams.forEach((value, key) => {
-      // Don't override base URL params
-      if (!searchParams.has(key)) {
-        searchParams.set(key, value);
-      }
-    });
-
-    searchParams.set(KnownQueryParams.sessionId, sessionId);
-    return searchParams;
+  /** Keep notebook page parameters separate from connection options. */
+  private getSessionURL(path: string, sessionId: SessionId): URL {
+    const url = this.formatHttpURL({ path });
+    const notebookParams = new URLSearchParams(this.httpURL.search);
+    const pageParams = new URLSearchParams(window.location.search);
+    for (const key of pageParams.keys()) {
+      notebookParams.delete(key);
+    }
+    for (const [key, value] of pageParams) {
+      notebookParams.append(key, value);
+    }
+    // These parameters are internal and are not exposed by mo.query_params().
+    notebookParams.delete(KnownQueryParams.accessToken);
+    notebookParams.delete("refresh_token");
+    notebookParams.delete(KnownQueryParams.sessionId);
+    url.searchParams.set("query_params", notebookParams.toString());
+    url.searchParams.set(KnownQueryParams.sessionId, sessionId);
+    return url;
   }
 
   /**
    * The WebSocket URL of the runtime.
    */
   public getWsURL(sessionId: SessionId): URL {
-    const searchParams = this.getSessionSearchParams(sessionId);
-    // TODO(MO-8039): Rediscover hosts and resolve notebooks on reconnect.
-    const notebookId = this.config.notebookId;
-    const path = notebookId
-      ? `/api/marimo/v1/notebooks/${encodeURIComponent(notebookId)}/events/ws`
-      : "/ws";
-    const url = this.formatWsURL(path, searchParams);
-    if (notebookId) {
-      // TODO(MO-8034): Separate host options from notebook query params.
-      url.searchParams.set("types", NOTEBOOK_EVENT_TYPES.join(","));
-    }
-    return url;
+    const url = this.getSessionURL("/ws", sessionId);
+    return asWsUrl(this.appendCrossOriginAuth(url).toString());
   }
 
   /**
@@ -176,13 +164,7 @@ export class RuntimeManager {
    * (see `headers()`) instead of the URL.
    */
   public getSseURL(sessionId: SessionId): URL {
-    // TODO: Migrate editor SSE to the host event stream. Until then, host
-    // state such as attachments is unavailable on this legacy connection.
-    const url = this.formatHttpURL({
-      path: "/sse",
-      searchParams: this.getSessionSearchParams(sessionId),
-      restrictToKnownQueryParams: false,
-    });
+    const url = this.getSessionURL("/sse", sessionId);
     // The page URL can still carry access_token (merged in from the
     // current page's params, before auth cleanup strips it); never
     // forward it — a token in the URL leaks into proxy/server logs.
@@ -191,10 +173,15 @@ export class RuntimeManager {
   }
 
   /**
-   * The WebSocket Sync URL of the runtime, for real-time updates.
+   * The host API URL of a notebook, below which its operations are served.
+   * Requests send auth in headers (see `headers()`), never in the URL.
    */
-  public getWsSyncURL(sessionId: SessionId): URL {
-    return this.formatWsURL("/ws_sync", this.getSessionSearchParams(sessionId));
+  public getNotebookURL(notebookId: NotebookId): URL {
+    const url = this.formatHttpURL({
+      path: `/api/marimo/v1/notebooks/${encodeURIComponent(notebookId)}`,
+    });
+    url.searchParams.delete(KnownQueryParams.accessToken);
+    return url;
   }
 
   /**

package/src/core/runtime/types.ts

@@ -1,8 +1,6 @@
 /* Copyright 2026 Marimo. All rights reserved. */
 
 export interface RuntimeConfig {
-  /** The host notebook for WebSocket events. SSE uses the legacy session stream. */
-  readonly notebookId?: string;
   /**
    * The URL of the runtime server.
    */

package/src/core/saving/__tests__/document.test.tsx

@@ -0,0 +1,123 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+import type * as Protocol from "@marimo-team/client/protocol";
+import { renderHook } from "@testing-library/react";
+import { atom, createStore } from "jotai";
+import type { PropsWithChildren } from "react";
+import { describe, expect, it, vi } from "vitest";
+import { DocumentWrapper } from "@/__mocks__/document";
+import { partial } from "@/__mocks__/partial";
+import type {
+  DocumentModel,
+  DocumentSyncStatus,
+} from "@/core/notebook/document";
+import type { NotebookReplica } from "@/core/notebook/replica";
+import {
+  type FakeSave,
+  fakeBackend,
+} from "@/core/notebook/replica/__tests__/fixtures";
+import { activeNotebookReplicaAtom } from "@/core/session/active-notebook";
+import { Deferred } from "@/utils/Deferred";
+import { useNotebookDocumentSave } from "../document";
+
+const confirmed: Protocol.Document = {
+  notebook_id: "n1",
+  order: [],
+  app_config: {},
+  etag: "e1",
+  saved_etag: null,
+  order_etag: "o1",
+  app_config_etag: "a1",
+  version: 1,
+};
+
+const saved: FakeSave = async () => ({
+  ok: true,
+  status: 200,
+  data: { document: confirmed, cells: {} },
+});
+
+function setup({
+  flush = vi.fn(async () => confirmed),
+  save = vi.fn(saved),
+  connected = true,
+}: {
+  flush?: DocumentModel["flush"];
+  save?: FakeSave;
+  connected?: boolean;
+} = {}) {
+  const store = createStore();
+  const model = partial<DocumentModel>(
+    {
+      flush,
+      retry: vi.fn(),
+      syncStatusAtom: atom<DocumentSyncStatus>({ status: "idle" }),
+    },
+    "model",
+  );
+  if (connected) {
+    store.set(
+      activeNotebookReplicaAtom,
+      partial<NotebookReplica>(
+        { backendAtom: atom(fakeBackend({ save })) },
+        "replica",
+      ),
+    );
+  }
+  const wrapper = ({ children }: PropsWithChildren) => (
+    <DocumentWrapper store={store} model={model}>
+      {children}
+    </DocumentWrapper>
+  );
+  const { result } = renderHook(() => useNotebookDocumentSave(), { wrapper });
+  return { saveDocument: result.current, flush, save };
+}
+
+describe("useNotebookDocumentSave", () => {
+  it("flushes this client's edits, then saves the flushed notebook", async () => {
+    const flushed = new Deferred<Protocol.Document>();
+    const save = vi.fn(saved);
+    const { saveDocument } = setup({ flush: () => flushed.promise, save });
+    const saving = saveDocument();
+    await Promise.resolve();
+    expect(save).not.toHaveBeenCalled();
+    flushed.resolve(confirmed);
+    await saving;
+    expect(save).toHaveBeenCalledWith("n1");
+  });
+
+  it("does not save when flushing fails", async () => {
+    const save = vi.fn(saved);
+    const { saveDocument } = setup({
+      flush: () => Promise.reject(new Error("unsynchronized edit")),
+      save,
+    });
+    await expect(saveDocument()).rejects.toThrow("unsynchronized edit");
+    expect(save).not.toHaveBeenCalled();
+  });
+
+  it("reports the host's refusal", async () => {
+    const { saveDocument } = setup({
+      save: async () => ({
+        ok: false,
+        status: 409,
+        problem: {
+          type: "about:blank",
+          title: "Conflict",
+          status: 409,
+          detail: "storage full",
+        },
+      }),
+    });
+    await expect(saveDocument()).rejects.toThrow("storage full");
+  });
+
+  it("does not save while no notebook is connected", async () => {
+    const flush = vi.fn(async () => confirmed);
+    const { saveDocument } = setup({ flush, connected: false });
+    await expect(saveDocument()).rejects.toThrow(
+      "The notebook is not connected",
+    );
+    expect(flush).not.toHaveBeenCalled();
+  });
+});

package/src/core/saving/__tests__/state.test.ts

@@ -0,0 +1,111 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+import type * as Protocol from "@marimo-team/client/protocol";
+import { atom, createStore } from "jotai";
+import { describe, expect, it } from "vitest";
+import { partial } from "@/__mocks__/partial";
+import type {
+  DocumentModel,
+  DocumentSyncStatus,
+} from "@/core/notebook/document";
+import type { Document, NotebookReplica } from "@/core/notebook/replica";
+import {
+  activeDocumentModelAtom,
+  activeNotebookReplicaAtom,
+} from "@/core/session/active-notebook";
+import { initialLayoutState, layoutStateAtom } from "../../layout/state";
+import { documentVersionAtom, needsSaveAtom, savedLayoutAtom } from "../state";
+
+function revision(
+  overrides: Partial<Protocol.Document> = {},
+): Protocol.Document {
+  return {
+    notebook_id: "nb",
+    order: [],
+    app_config: {},
+    etag: "v2",
+    saved_etag: "v2",
+    order_etag: "order",
+    app_config_etag: "config",
+    version: 2,
+    ...overrides,
+  };
+}
+
+function connect(
+  store: ReturnType<typeof createStore>,
+  document: Document | null,
+  status: DocumentSyncStatus["status"] = "idle",
+) {
+  store.set(
+    activeNotebookReplicaAtom,
+    partial<NotebookReplica>({ documentAtom: atom(document) }, "replica"),
+  );
+  store.set(
+    activeDocumentModelAtom,
+    partial<DocumentModel>(
+      { syncStatusAtom: atom({ status } as DocumentSyncStatus) },
+      "model",
+    ),
+  );
+}
+
+describe("needsSaveAtom", () => {
+  it("is false without a connected document", () => {
+    const store = createStore();
+    expect(store.get(needsSaveAtom)).toBe(false);
+    expect(store.get(documentVersionAtom)).toBeNull();
+    connect(store, null);
+    expect(store.get(needsSaveAtom)).toBe(false);
+  });
+
+  it("compares the host's etag with what it last wrote", () => {
+    const store = createStore();
+    connect(store, { revision: revision(), cells: new Map() });
+    expect(store.get(needsSaveAtom)).toBe(false);
+    expect(store.get(documentVersionAtom)).toBe(2);
+    connect(store, { revision: revision({ etag: "v3" }), cells: new Map() });
+    expect(store.get(needsSaveAtom)).toBe(true);
+  });
+
+  it("treats a document never written as unsaved", () => {
+    const store = createStore();
+    connect(store, {
+      revision: revision({ saved_etag: null }),
+      cells: new Map(),
+    });
+    expect(store.get(needsSaveAtom)).toBe(true);
+  });
+
+  it.each(["pending", "sending", "confirming", "error"] as const)(
+    "counts edits the host has not confirmed (%s)",
+    (status) => {
+      const store = createStore();
+      connect(store, { revision: revision(), cells: new Map() }, status);
+      expect(store.get(needsSaveAtom)).toBe(true);
+    },
+  );
+
+  it("ignores a model that is not running", () => {
+    const store = createStore();
+    connect(store, { revision: revision(), cells: new Map() }, "unavailable");
+    expect(store.get(needsSaveAtom)).toBe(false);
+  });
+
+  it("compares the layout with the one last saved", () => {
+    const store = createStore();
+    connect(store, { revision: revision(), cells: new Map() });
+    store.set(layoutStateAtom, {
+      selectedLayout: "grid",
+      layoutData: { grid: null },
+    });
+    expect(store.get(needsSaveAtom)).toBe(false);
+    store.set(savedLayoutAtom, initialLayoutState());
+    expect(store.get(needsSaveAtom)).toBe(true);
+    store.set(savedLayoutAtom, {
+      selectedLayout: "grid",
+      layoutData: { grid: null },
+    });
+    expect(store.get(needsSaveAtom)).toBe(false);
+  });
+});

package/src/core/saving/__tests__/useAutoSave.test.tsx

@@ -0,0 +1,67 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+import { act, renderHook } from "@testing-library/react";
+import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
+import { ConnectionState } from "../../session/types";
+import { useAutoSave } from "../useAutoSave";
+
+function options(): Parameters<typeof useAutoSave>[0] {
+  return {
+    version: 1,
+    config: {
+      autosave: "after_delay",
+      autosave_delay: 1000,
+      format_on_save: false,
+    },
+    connStatus: { state: ConnectionState.OPEN },
+    needsSave: true,
+    canEdit: true,
+    onSave: vi.fn(),
+  };
+}
+
+describe("useAutoSave", () => {
+  beforeEach(() => {
+    vi.useFakeTimers();
+  });
+
+  afterEach(() => {
+    vi.useRealTimers();
+  });
+
+  it("waits for the full delay after the latest confirmed change", () => {
+    const initial = options();
+    const { rerender } = renderHook(useAutoSave, { initialProps: initial });
+
+    act(() => vi.advanceTimersByTime(600));
+    rerender({ ...initial, version: 2 });
+    act(() => vi.advanceTimersByTime(999));
+    expect(initial.onSave).not.toHaveBeenCalled();
+
+    act(() => vi.advanceTimersByTime(1));
+    expect(initial.onSave).toHaveBeenCalledExactlyOnceWith();
+  });
+
+  it("cancels pending saves when editing is disabled or the component unmounts", () => {
+    const initial = options();
+    const { rerender, unmount } = renderHook(useAutoSave, {
+      initialProps: initial,
+    });
+
+    act(() => vi.advanceTimersByTime(600));
+    rerender({ ...initial, canEdit: false });
+    act(() => vi.advanceTimersByTime(1000));
+    expect(initial.onSave).not.toHaveBeenCalled();
+
+    rerender(initial);
+    act(() => vi.advanceTimersByTime(999));
+    expect(initial.onSave).not.toHaveBeenCalled();
+    act(() => vi.advanceTimersByTime(1));
+    expect(initial.onSave).toHaveBeenCalledExactlyOnceWith();
+
+    rerender({ ...initial, version: 2 });
+    unmount();
+    act(() => vi.advanceTimersByTime(1000));
+    expect(initial.onSave).toHaveBeenCalledExactlyOnceWith();
+  });
+});