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

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

@@ -13,7 +13,9 @@ describe("RuntimeManager.getWsURL", () => {
     });
     const sessionId = "1234" as SessionId;
     const url = runtime.getWsURL(sessionId);
-    expect(url.toString()).toBe("wss://marimo.app/ws?session_id=1234");
+    expect(url.toString()).toBe(
+      "wss://marimo.app/ws?query_params=&session_id=1234",
+    );
     expect(url.searchParams.get(KnownQueryParams.sessionId)).toBe(sessionId);
   });
 
@@ -24,7 +26,9 @@ describe("RuntimeManager.getWsURL", () => {
     });
     const sessionId = "1234" as SessionId;
     const url = runtime.getWsURL(sessionId);
-    expect(url.toString()).toBe("ws://marimo.app/ws?session_id=1234");
+    expect(url.toString()).toBe(
+      "ws://marimo.app/ws?query_params=&session_id=1234",
+    );
     expect(url.searchParams.get(KnownQueryParams.sessionId)).toBe(sessionId);
   });
 
@@ -35,7 +39,9 @@ describe("RuntimeManager.getWsURL", () => {
     });
     const sessionId = "1234" as SessionId;
     const url = runtime.getWsURL(sessionId);
-    expect(url.toString()).toBe("ws://marimo.app/nested/ws?session_id=1234");
+    expect(url.toString()).toBe(
+      "ws://marimo.app/nested/ws?query_params=&session_id=1234",
+    );
     expect(url.searchParams.get(KnownQueryParams.sessionId)).toBe(sessionId);
   });
 
@@ -48,7 +54,7 @@ describe("RuntimeManager.getWsURL", () => {
     const sessionId = "1234" as SessionId;
     const url = runtime.getWsURL(sessionId);
     expect(url.toString()).toBe(
-      "ws://marimo.app/nested/ws?foo=bar&file=test.py&session_id=1234",
+      "ws://marimo.app/nested/ws?foo=bar&file=test.py&query_params=foo%3Dbar%26file%3Dtest.py&session_id=1234",
     );
     expect(url.searchParams.get(KnownQueryParams.sessionId)).toBe(sessionId);
   });
@@ -62,11 +68,75 @@ describe("RuntimeManager.getWsURL", () => {
     const sessionId = "1234" as SessionId;
     const url = runtime.getWsURL(sessionId);
     expect(url.toString()).toBe(
-      "ws://marimo.app/ws?file=test.py&search=test&other=last&session_id=1234",
+      "ws://marimo.app/ws?file=test.py&query_params=file%3Dtest.py%26search%3Dtest%26other%3Dlast&session_id=1234",
     );
     expect(url.searchParams.get("file")).toBe("test.py");
-    expect(url.searchParams.get("search")).toBe("test");
-    expect(url.searchParams.get("other")).toBe("last");
+    expect(url.searchParams.get("query_params")).toBe(
+      "file=test.py&search=test&other=last",
+    );
     expect(url.searchParams.get(KnownQueryParams.sessionId)).toBe(sessionId);
   });
 });
+
+describe.each(["getWsURL", "getSseURL"] as const)(
+  "RuntimeManager.%s",
+  (method) => {
+    it("preserves runtime parameters and repeated page values", () => {
+      window.history.pushState({}, "", "/?tag=page-one&tag=page-two");
+      const runtime = new RuntimeManager(
+        {
+          url: `${window.location.origin}/nested?proxy_token=secret&tag=base&base=one&base=two`,
+          lazy: true,
+        },
+        true,
+      );
+      const url = runtime[method]("session" as SessionId);
+      expect([...url.searchParams]).toEqual([
+        ["proxy_token", "secret"],
+        ["tag", "base"],
+        ["base", "one"],
+        ["base", "two"],
+        [
+          "query_params",
+          "proxy_token=secret&base=one&base=two&tag=page-one&tag=page-two",
+        ],
+        ["session_id", "session"],
+      ]);
+    });
+
+    it("keeps notebook parameters out of the stream options", () => {
+      const notebookParams = new URLSearchParams([
+        ["types", "scatter"],
+        ["cursor", "page-2"],
+        ["run_id", "experiment"],
+        ["query_params", "nested=value&other=1"],
+        ["tag", "one"],
+        ["tag", "two"],
+        ["empty", ""],
+        ["text", "a+b & café"],
+        ["file", "test.py"],
+        ["kiosk", "true"],
+      ]);
+      window.history.pushState(
+        {},
+        "",
+        `/?${notebookParams}&session_id=page-session`,
+      );
+      const runtime = new RuntimeManager(
+        { url: window.location.origin, lazy: true },
+        true,
+      );
+      const url = runtime[method]("actual-session" as SessionId);
+
+      expect(Object.fromEntries(url.searchParams)).toEqual({
+        file: "test.py",
+        kiosk: "true",
+        session_id: "actual-session",
+        query_params: notebookParams.toString(),
+      });
+      expect([
+        ...new URLSearchParams(url.searchParams.get("query_params") ?? ""),
+      ]).toEqual([...notebookParams]);
+    });
+  },
+);

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

@@ -8,8 +8,9 @@ import { RuntimeManager } from "../runtime";
 import type { RuntimeConfig } from "../types";
 
 // Mock the session module
-vi.mock("@/core/kernel/session", () => ({
-  getSessionId: () => "test-session-id" as SessionId,
+vi.mock("@/core/kernel/session", async (importOriginal) => ({
+  ...(await importOriginal<typeof import("@/core/kernel/session")>()),
+  getSessionId: () => "test-session-id",
 }));
 
 vi.mock("@/utils/Logger", () => ({ Logger: Mocks.quietLogger() }));
@@ -59,6 +60,7 @@ describe("RuntimeManager", () => {
       expect(url.protocol).toBe("wss:");
       expect(url.hostname).toBe("example.com");
       expect(url.pathname).toBe("/ws");
+      expect(url.searchParams.has("types")).toBe(false);
       expect(url.searchParams.get("session_id")).toBe(sessionId);
     });
 
@@ -218,16 +220,24 @@ describe("RuntimeManager", () => {
     });
   });
 
-  describe("getSseURL", () => {
-    it("should return an http(s) URL with the session ID", () => {
-      const runtime = new RuntimeManager(mockConfig);
-      const url = runtime.getSseURL("1234" as 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.protocol).toBe("https:");
-      expect(url.pathname).toBe("/sse");
-      expect(url.searchParams.get("session_id")).toBe("1234");
+      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
@@ -256,7 +266,8 @@ 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);
       }
@@ -676,9 +687,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");
@@ -737,7 +755,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,
@@ -745,7 +763,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

@@ -8,7 +8,7 @@ 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";
 

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,5 +1,6 @@
 /* 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";
@@ -9,6 +10,8 @@ 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;
@@ -71,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);
@@ -129,32 +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 {
-    return this.formatWsURL("/ws", this.getSessionSearchParams(sessionId));
+    const url = this.getSessionURL("/ws", sessionId);
+    return asWsUrl(this.appendCrossOriginAuth(url).toString());
   }
 
   /**
@@ -164,11 +164,7 @@ export class RuntimeManager {
    * (see `headers()`) instead of the URL.
    */
   public getSseURL(sessionId: SessionId): URL {
-    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.
@@ -176,11 +172,24 @@ export class RuntimeManager {
     return url;
   }
 
+  /**
+   * 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 getNotebookURL(notebookId: NotebookId): URL {
+    const url = this.formatHttpURL({
+      path: `/api/marimo/v1/notebooks/${encodeURIComponent(notebookId)}`,
+    });
+    url.searchParams.delete(KnownQueryParams.accessToken);
+    return url;
+  }
+
   /**
    * The WebSocket Sync URL of the runtime, for real-time updates.
    */
   public getWsSyncURL(sessionId: SessionId): URL {
-    return this.formatWsURL("/ws_sync", this.getSessionSearchParams(sessionId));
+    const url = this.getSessionURL("/ws_sync", sessionId);
+    return asWsUrl(this.appendCrossOriginAuth(url).toString());
   }
 
   /**

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

package/src/core/saving/filename.ts

@@ -9,7 +9,7 @@ import { getAppConfig } from "../config/config";
 import { KnownQueryParams } from "../constants";
 import { connectionAtom } from "../network/connection";
 import { useRequestClient } from "../network/requests";
-import { ConnectionState } from "../websocket/types";
+import { ConnectionState } from "../session/types";
 import { filenameAtom } from "./file-state";
 
 export function useFilename() {