@marimo-team/islands 0.25.3-dev31 → 0.25.3-dev311

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);
@@ -558,6 +509,42 @@ describe("RuntimeManager", () => {
     });
   });
 
+  describe("start", () => {
+    it("starts once and lets later callers join the same attempt", async () => {
+      const runtime = new RuntimeManager(mockConfig, true);
+      const healthySpy = vi
+        .spyOn(runtime, "reconcileFromHealth")
+        .mockResolvedValue(true);
+
+      await Promise.all([runtime.start(), runtime.start()]);
+
+      expect(healthySpy).toHaveBeenCalledTimes(1);
+    });
+
+    it("tries again after an attempt gave up", async () => {
+      const runtime = new RuntimeManager(mockConfig, true);
+      const healthySpy = vi
+        .spyOn(runtime, "reconcileFromHealth")
+        .mockResolvedValue(false);
+      vi.spyOn(runtime, "init").mockImplementation(async () => {
+        await runtime.reconcileFromHealth();
+        throw new Error("Failed to connect after 25 retries");
+      });
+
+      await expect(runtime.start()).rejects.toThrow("Failed to connect");
+      healthySpy.mockResolvedValue(true);
+      vi.spyOn(runtime, "init").mockImplementation(async () => {
+        await runtime.reconcileFromHealth();
+        await runtime.waitForHealthy().catch(() => undefined);
+      });
+
+      // The second attempt is a new one, not the old failure replayed.
+      vi.spyOn(runtime, "waitForHealthy").mockResolvedValue(undefined);
+      await expect(runtime.start()).resolves.toBeUndefined();
+      expect(healthySpy).toHaveBeenCalledTimes(2);
+    });
+  });
+
   describe("headers", () => {
     it("should return headers with auth token", () => {
       const runtime = new RuntimeManager(mockConfig);
@@ -720,9 +707,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 +775,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 +783,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

@@ -1,11 +1,13 @@
 /* Copyright 2026 Marimo. All rights reserved. */
 import { atom, useAtom, useAtomValue } from "jotai";
 import useEvent from "react-use-event-hook";
+import { prettyError } from "@/utils/errors";
 import { Logger } from "@/utils/Logger";
 import { connectionAtom } from "../network/connection";
+import { ConnectionClosedReason } from "../notebook/replica";
 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";
 
@@ -37,7 +39,7 @@ export function useConnectToRuntime(): () => Promise<void> {
   return useEvent(async () => {
     if (isAppNotStarted(connection.state)) {
       setConnection({ state: ConnectionState.CONNECTING });
-      await runtimeManager.init();
+      await runtimeManager.start();
     } else {
       Logger.log("Runtime already started or starting...");
     }
@@ -51,6 +53,39 @@ export function getRuntimeManager(): RuntimeManager {
   return store.get(runtimeManagerAtom);
 }
 
+/**
+ * Start a lazy runtime the first time something needs it.
+ *
+ * The connection can already read as connecting before anyone has asked
+ * for the runtime, because the notebook stream waits for the runtime's
+ * health and the editor mirrors that wait. So whether the runtime has
+ * been asked to start is the runtime's own business, and this only keeps
+ * the connection state honest around the attempt. A start that gives up
+ * closes the connection with the reason, which leaves the user a way to
+ * reconnect instead of a connection that looks like it is still coming up.
+ */
+export async function startRuntimeIfLazy(): Promise<void> {
+  const runtimeManager = getRuntimeManager();
+  if (runtimeManager.isServerless || !runtimeManager.isLazy) {
+    return;
+  }
+  if (isAppNotStarted(store.get(connectionAtom).state)) {
+    store.set(connectionAtom, { state: ConnectionState.CONNECTING });
+  }
+  try {
+    await runtimeManager.start();
+  } catch (error) {
+    if (store.get(connectionAtom).state === ConnectionState.CONNECTING) {
+      store.set(connectionAtom, {
+        state: ConnectionState.CLOSED,
+        code: ConnectionClosedReason.KERNEL_STARTUP_ERROR,
+        reason: prettyError(error),
+      });
+    }
+    throw error;
+  }
+}
+
 export function asRemoteURL(path: string): URL {
   if (path.startsWith("http")) {
     return new URL(path);

package/src/core/runtime/runtime.ts

@@ -1,17 +1,20 @@
 /* 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 starting: Promise<void> | undefined;
   private config: RuntimeConfig;
   private lazy: boolean;
 
@@ -45,7 +48,7 @@ export class RuntimeManager {
     return this.httpURL.origin === window.location.origin;
   }
 
-  private get isServerless(): boolean {
+  public get isServerless(): boolean {
     return isWasm() || isIslands() || isStaticNotebook();
   }
 
@@ -72,10 +75,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 +130,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 +165,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 +174,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;
   }
 
   /**
@@ -358,6 +346,25 @@ export class RuntimeManager {
     return this.initialHealthyCheck.promise;
   }
 
+  /**
+   * Start the runtime once and wait until it is healthy.
+   *
+   * A run, the connect button, and the notebook stream can each ask for
+   * the runtime, so the first call starts it and later calls join that
+   * attempt. An attempt that gave up is forgotten, so the next call tries
+   * again instead of failing on the old result.
+   */
+  public start(): Promise<void> {
+    this.starting ??= this.init()
+      .then(() => this.waitForHealthy())
+      .catch((error: unknown) => {
+        this.starting = undefined;
+        this.initialHealthyCheck = new Deferred<void>();
+        throw error;
+      });
+    return this.starting;
+  }
+
   public headers(): KnownHeaders {
     const headers: KnownHeaders = {
       "Marimo-Session-Id": getSessionId(),

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