@marimo-team/islands 0.25.3-dev16 → 0.25.3-dev167

package/src/core/notebook/replica/replica.ts

@@ -0,0 +1,286 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+/** What a backend's stream has confirmed about a notebook, kept in atoms. */
+
+import { Logger } from "@/utils/Logger";
+import {
+  createNotebookReducer,
+  type NotebookEvent,
+  NotebookEvents,
+  StreamStopped,
+} from "@marimo-team/client";
+import type { Attachment, NotebookOutput } from "@marimo-team/client/protocol";
+import { type Atom, atom, type createStore } from "jotai/vanilla";
+import { z } from "zod";
+import type { NotebookBackend } from "./backend";
+import {
+  ConnectionClosedReason,
+  ConnectionState,
+  type ConnectionStatus,
+} from "./connection";
+import { attachments } from "./resources/attachments";
+import { type Document, document } from "./resources/document";
+
+/** The kinds the editor reads, with marimo's own notifications. */
+const TYPES = ["notebook", "cell", "document", "attachment", "marimo"] as const;
+
+type EditorEvent = NotebookEvent<(typeof TYPES)[number]>;
+
+/** A native marimo notification, passed through without interpretation. */
+const MarimoMessage = z.object({ marimo: z.unknown() });
+
+/** A notebook as its backend's stream reports it, kept current while connected. */
+export interface NotebookReplica {
+  /** Connection status, independent of application session readiness. */
+  readonly connectionAtom: Atom<ConnectionStatus>;
+  /** Identity and capabilities, available while the stream is connected and ready. */
+  readonly notebookAtom: Atom<NotebookOutput | null>;
+  /**
+   * The most recently confirmed document.
+   *
+   * Null until a snapshot is complete, after reset or removal, and on
+   * connections without host document events. A temporary disconnection
+   * retains the last document. Partial updates are not exposed.
+   */
+  readonly documentAtom: Atom<Document | null>;
+  /**
+   * The backend this replica follows, or null while not connected.
+   *
+   * Its commands reach the host that confirmed `documentAtom`. A temporary
+   * disconnection keeps it. Connecting again or disconnecting replaces it.
+   */
+  readonly backendAtom: Atom<NotebookBackend | null>;
+  /** Live attachments, empty until the stream is ready or while disconnected. */
+  readonly attachmentsAtom: Atom<readonly Attachment[]>;
+  /**
+   * Follow a backend, replacing any current connection.
+   *
+   * The returned function ends this connection and clears the replica.
+   * It does nothing once another connection has replaced this one.
+   */
+  connect(
+    backend: NotebookBackend,
+    options: {
+      /** Handle a native marimo notification. Its payload has not been validated. */
+      onNotification: (notification: unknown) => void;
+    },
+  ): () => void;
+  /** Restore the current connection after its stream closed, if the host is healthy. */
+  reconnect(): Promise<void>;
+}
+
+const disconnected = {
+  state: ConnectionState.CLOSED,
+  code: ConnectionClosedReason.KERNEL_DISCONNECTED,
+  reason: "kernel not found",
+} satisfies ConnectionStatus;
+
+const malformed = {
+  ...disconnected,
+  reason: "Failed to decode notebook message",
+} satisfies ConnectionStatus;
+
+function isStatus(value: unknown): value is ConnectionStatus {
+  return typeof value === "object" && value !== null && "state" in value;
+}
+
+/**
+ * The status to show after the stream stopped.
+ *
+ * A transport says which status its stop means. An event the protocol
+ * rejects is malformed. Anything else means the host was not reached.
+ */
+function statusAfter(error: unknown): ConnectionStatus {
+  if (error instanceof StreamStopped && isStatus(error.cause)) {
+    return error.cause;
+  }
+  if (error instanceof StreamStopped && error.cause instanceof z.ZodError) {
+    return malformed;
+  }
+  return disconnected;
+}
+
+/** Create a replica without connecting it. */
+export function createNotebookReplica(
+  store: ReturnType<typeof createStore>,
+): NotebookReplica {
+  // `useNotebookSession` copies this into the editor's shared `connectionAtom`,
+  // which gates features until MO-8051 derives availability separately.
+  const connectionAtom = atom<ConnectionStatus>({
+    state: ConnectionState.NOT_STARTED,
+  });
+  // One atom holds the whole stream, which resource atoms derive from.
+  const reducer = createNotebookReducer({ document, attachments });
+  const streamAtom = atom(reducer.initial);
+  // A resumed stream is live, but presence waits for this connection's ready.
+  const readyAtom = atom(false);
+  const presentAtom = atom(
+    (get) =>
+      get(readyAtom) && get(connectionAtom).state === ConnectionState.OPEN,
+  );
+  const notebookAtom = atom((get) =>
+    get(presentAtom) ? get(streamAtom).notebook : null,
+  );
+  const documentAtom = atom(
+    (get) => get(streamAtom).resources.document.current,
+  );
+  const attachmentsAtom = atom((get) =>
+    get(presentAtom)
+      ? get(streamAtom).resources.attachments
+      : attachments.initial,
+  );
+  const backendAtom = atom<NotebookBackend | null>(null);
+  /** The connection this replica follows. Connecting again ends it. */
+  let current: { end(): void; reconnect(): Promise<void> } | undefined;
+
+  function setStatus(status: ConnectionStatus) {
+    store.set(readyAtom, false);
+    store.set(connectionAtom, status);
+  }
+
+  /** Drop what the stream reported and show `status`. */
+  function clear(status: ConnectionStatus) {
+    store.set(streamAtom, reducer.initial);
+    setStatus(status);
+  }
+
+  return {
+    connectionAtom,
+    notebookAtom,
+    documentAtom,
+    backendAtom,
+    attachmentsAtom,
+    connect,
+    reconnect: async () => current?.reconnect(),
+  };
+
+  function connect(
+    connected: NotebookBackend,
+    { onNotification }: { onNotification: (notification: unknown) => void },
+  ) {
+    current?.end();
+    const lifetime = new AbortController();
+    const { signal } = lifetime;
+    const checkHealth = connected.checkHealth ?? (async () => true);
+    store.set(backendAtom, connected);
+    const { transport } = connected;
+    let events: NotebookEvents<(typeof TYPES)[number]> | undefined;
+    let checkingHealth = false;
+    // Whether the replica still holds the state the last cursor describes.
+    let resumable = true;
+
+    /** Mirror one stream into the atoms until it closes. */
+    function follow(cursor?: string) {
+      if (cursor === undefined) {
+        store.set(streamAtom, reducer.initial);
+      }
+      resumable = true;
+      const current = new NotebookEvents(transport, {
+        types: TYPES,
+        cursor,
+        signal,
+        // TODO(MO-8051): Gate editing and running on the notebook's operations
+        // and request results instead. The stream's status would then only
+        // mark the view as stale.
+        onStatus: (status) => {
+          if (status === "connecting") {
+            setStatus({ state: ConnectionState.CONNECTING });
+          } else if (status === "connected") {
+            setStatus({ state: ConnectionState.OPEN });
+          }
+        },
+      });
+      current.subscribe(
+        (event) => apply(current, event),
+        (error) => {
+          if (!signal.aborted && current === events) {
+            Logger.warn("Notebook stream stopped", error);
+            const status = statusAfter(error);
+            if (status === malformed) {
+              stopMalformed(current);
+            } else {
+              setStatus(status);
+            }
+          }
+        },
+      );
+      events = current;
+    }
+
+    function stopMalformed(source: NotebookEvents<string>) {
+      source.close();
+      // The cursor can already include an event the replica rejected.
+      resumable = false;
+      clear(malformed);
+    }
+
+    function apply(source: NotebookEvents<string>, event: EditorEvent) {
+      if (event.event === "marimo") {
+        const message = MarimoMessage.safeParse(event.data);
+        if (message.success) {
+          onNotification(message.data.marimo);
+        } else {
+          Logger.error("Failed to decode notebook message", event);
+          stopMalformed(source);
+        }
+        return;
+      }
+      try {
+        store.set(streamAtom, reducer.reduce(store.get(streamAtom), event));
+      } catch (error) {
+        // An event the replica cannot apply, such as another notebook's.
+        Logger.error("Failed to apply notebook event", error);
+        stopMalformed(source);
+        return;
+      }
+      if (event.event === "ready" || event.event === "reset") {
+        store.set(readyAtom, event.event === "ready");
+      }
+    }
+
+    const reconnect = async () => {
+      if (signal.aborted || checkingHealth || events?.status !== "closed") {
+        return;
+      }
+      checkingHealth = true;
+      setStatus({ state: ConnectionState.CONNECTING });
+      try {
+        const healthy = await checkHealth();
+        if (signal.aborted) {
+          return;
+        }
+        if (!healthy) {
+          setStatus(disconnected);
+          return;
+        }
+        follow(resumable ? (events.cursor ?? undefined) : undefined);
+      } catch (error) {
+        if (!signal.aborted) {
+          Logger.error("Notebook health check failed", error);
+          setStatus(disconnected);
+        }
+      } finally {
+        checkingHealth = false;
+      }
+    };
+
+    const connection = {
+      end: () => {
+        if (signal.aborted) {
+          return;
+        }
+        lifetime.abort();
+        if (current === connection) {
+          current = undefined;
+          store.set(backendAtom, null);
+        }
+        clear({ state: ConnectionState.NOT_STARTED });
+      },
+      reconnect,
+    };
+    current = connection;
+    setStatus({ state: ConnectionState.CONNECTING });
+    follow();
+    return connection.end;
+  }
+}

package/src/core/notebook/replica/resources/attachments.ts

@@ -0,0 +1,22 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+import type { Attachment } from "@marimo-team/client/protocol";
+import type { NotebookResource } from "@marimo-team/client";
+
+/** Attachment records retained with the stream's resume cursor. */
+export const attachments: NotebookResource<
+  readonly Attachment[],
+  "attachment"
+> = {
+  initial: [],
+  on: {
+    attachment: (state, { data: { attachment } }) =>
+      state.some((item) => item.id === attachment.id)
+        ? state.map((item) => (item.id === attachment.id ? attachment : item))
+        : [...state, attachment],
+    "attachment.removed": (state, { data }) =>
+      state.some((item) => item.id === data.id)
+        ? state.filter((item) => item.id !== data.id)
+        : state,
+  },
+};

package/src/core/notebook/replica/resources/document.ts

@@ -0,0 +1,101 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+import type * as Protocol from "@marimo-team/client/protocol";
+import type { NotebookResource } from "@marimo-team/client";
+
+/**
+ * The notebook's content at one revision.
+ *
+ * The revision record carries the cell order, app config, etags, and
+ * version. Each cell in its order has a record in `cells`.
+ */
+export interface Document {
+  readonly revision: Protocol.Document;
+  readonly cells: ReadonlyMap<Protocol.CellId, Protocol.Cell>;
+}
+
+/** The records of a revision still being received. */
+interface PendingDocument {
+  /** Null until the `document` event that closes the group arrives. */
+  readonly revision: Protocol.Document | null;
+  readonly cells: ReadonlyMap<Protocol.CellId, Protocol.Cell>;
+}
+
+/**
+ * The confirmed document and the revision being assembled after it.
+ *
+ * A snapshot's records are pending until `ready`. A live revision is pending
+ * from its first cell until its `document` event.
+ */
+interface DocumentState {
+  readonly current: Document | null;
+  readonly pending: PendingDocument | null;
+}
+
+function pendingOf(state: DocumentState): PendingDocument {
+  return (
+    state.pending ?? {
+      revision: state.current?.revision ?? null,
+      cells: state.current?.cells ?? new Map(),
+    }
+  );
+}
+
+/** Confirm a revision whose records are complete. */
+function publish(
+  revision: Protocol.Document,
+  cells: ReadonlyMap<Protocol.CellId, Protocol.Cell>,
+): DocumentState {
+  const order = new Set(revision.order);
+  if (
+    order.size !== revision.order.length ||
+    order.size !== cells.size ||
+    [...order].some((id) => !cells.has(id)) ||
+    [...cells.values()].some((cell) => cell.version > revision.version)
+  ) {
+    throw new Error("Document records do not match its revision and order");
+  }
+  return { current: { revision, cells }, pending: null };
+}
+
+export const document: NotebookResource<DocumentState, "cell" | "document"> = {
+  initial: { current: null, pending: null },
+  on: {
+    cell: (state, { data }) => {
+      const pending = pendingOf(state);
+      return {
+        ...state,
+        pending: {
+          revision: pending.revision,
+          cells: new Map(pending.cells).set(data.cell.id, data.cell),
+        },
+      };
+    },
+    "cell.removed": (state, { data }) => {
+      const pending = pendingOf(state);
+      if (!pending.cells.has(data.id)) {
+        return state;
+      }
+      const cells = new Map(pending.cells);
+      cells.delete(data.id);
+      return { ...state, pending: { ...pending, cells } };
+    },
+    document: (state, { data }, { phase }) => {
+      const { cells } = pendingOf(state);
+      return phase === "snapshot"
+        ? { ...state, pending: { revision: data.document, cells } }
+        : publish(data.document, cells);
+    },
+    ready: (state, _event, { phase }) => {
+      if (state.pending === null) {
+        return state;
+      }
+      if (phase === "live") {
+        throw new Error("Stream ended with an incomplete document update");
+      }
+      return state.pending.revision === null
+        ? { current: null, pending: null }
+        : publish(state.pending.revision, state.pending.cells);
+    },
+  },
+};

package/src/core/packages/__tests__/useSandboxController.test.tsx

@@ -17,7 +17,7 @@ import type {
 } from "@/core/network/types";
 import { filenameAtom } from "@/core/saving/file-state";
 import { store } from "@/core/state/jotai";
-import { ConnectionState } from "@/core/websocket/types";
+import { ConnectionState } from "@/core/session/types";
 import { HTTPError } from "@/utils/errors";
 import { packageDataVersionAtom } from "../package-data";
 import {

package/src/core/packages/useSandboxController.ts

@@ -8,7 +8,7 @@ import { useRequestClient } from "@/core/network/requests";
 import { useFilename } from "@/core/saving/filename";
 import { waitFor } from "@/core/state/jotai";
 import { isWasm } from "@/core/wasm/utils";
-import { ConnectionState } from "@/core/websocket/types";
+import { ConnectionState } from "@/core/session/types";
 import { HTTPError, prettyError } from "@/utils/errors";
 import { invalidatePackageData } from "./package-data";
 import {

package/src/core/run-app.tsx

@@ -21,8 +21,8 @@ import { RuntimeState } from "./kernel/RuntimeState";
 import { getSessionId } from "./kernel/session";
 import { connectionAtom } from "./network/connection";
 import { useRequestClient } from "./network/requests";
-import { isAppConnecting } from "./websocket/connection-utils";
-import { useNotebookSession } from "./websocket/useNotebookSession";
+import { isAppConnecting } from "./session/connection-utils";
+import { useNotebookSession } from "./session/useNotebookSession";
 
 interface AppProps {
   appConfig: AppConfig;
@@ -49,7 +49,7 @@ export const RunApp: React.FC<AppProps> = ({ appConfig, hideHeader }) => {
     return () => {
       RuntimeState.INSTANCE.stop();
     };
-  }, []);
+  }, [sendComponentValues]);
 
   const { connection, reconnect } = useNotebookSession({
     autoInstantiate: true,

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

@@ -6,7 +6,7 @@ import { islandsInitializedAtom } from "../../islands/state";
 import { kernelStateAtom } from "../../kernel/state";
 import { connectionAtom } from "../../network/connection";
 import { wasmInitStateAtom } from "../../wasm/state";
-import { ConnectionState } from "../../websocket/types";
+import { ConnectionState } from "../../session/types";
 import {
   canRunCellsAtom,
   islandsAdapter,

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