@marimo-team/islands 0.25.3-dev17 → 0.25.3-dev173

package/src/core/session/backend.ts

@@ -0,0 +1,125 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+/** The backend the editor follows, chosen once for its mode. */
+
+import { NotebookClient, StreamStopped } from "@marimo-team/client";
+import type { NotebookId } from "@marimo-team/client/protocol";
+import type { NotebookBackend } from "../notebook/replica";
+import { PyodideBridge } from "../wasm/bridge";
+import { isWasm } from "../wasm/utils";
+import { sessionSse, websocket, withHeaders } from "../transports/server";
+import { idle } from "../transports/static";
+
+/** How the editor follows its notebook. */
+export interface EditorBackendOptions {
+  static: boolean;
+  transportType: "websocket" | "sse";
+  /** The session route, read before each connection. */
+  url: () => string;
+  /** The runtime's headers, read before each request. */
+  headers: () => Record<string, string>;
+  /** The host URL of a notebook the stream reported. */
+  notebookUrl: (notebookId: NotebookId) => URL;
+  /** Wait until the host can accept the first connection. */
+  waitUntilReady: () => Promise<void>;
+  /** Whether the host is reachable before an explicit reconnect. */
+  checkHealth: () => Promise<boolean>;
+}
+
+/**
+ * Read the session route once the host has started.
+ *
+ * Every connection waits for the same startup. Once it settles, even by
+ * giving up, connections read the route directly. An explicit reconnect
+ * checks the host's health instead.
+ */
+function afterStartup({
+  url,
+  waitUntilReady,
+}: Pick<EditorBackendOptions, "url" | "waitUntilReady">) {
+  let startup: Promise<void> | undefined;
+  let settled = false;
+  return async () => {
+    if (!settled) {
+      startup ??= waitUntilReady().finally(() => {
+        settled = true;
+      });
+      try {
+        await startup;
+      } catch (error) {
+        throw new StreamStopped("Healthy connection never made", {
+          cause: error,
+        });
+      }
+    }
+    return url();
+  };
+}
+
+/**
+ * Follow the notebook running in the WASM worker.
+ *
+ * The worker's backend exists once the worker has started, so the first
+ * connection creates it. Commands are only for a notebook its stream reported.
+ */
+function wasm(): NotebookBackend {
+  let backend: Promise<NotebookBackend> | undefined;
+  let started: NotebookBackend | undefined;
+  return {
+    transport: async (attempt, handlers) => {
+      backend ??= PyodideBridge.INSTANCE.createNotebookBackend();
+      try {
+        started = await backend;
+      } catch (error) {
+        throw new StreamStopped("The WASM notebook did not start", {
+          cause: error,
+        });
+      }
+      return started.transport(attempt, handlers);
+    },
+    notebook: (id) => {
+      if (started === undefined) {
+        throw new Error("The WASM notebook has not started");
+      }
+      return started.notebook(id);
+    },
+  };
+}
+
+/** Commands for a notebook the editor cannot change through its backend. */
+function uneditable(): ReturnType<NotebookBackend["notebook"]> {
+  const unsupported: NotebookClient["save"] = async () => ({
+    ok: false,
+    status: 501,
+    problem: {
+      type: "tag:marimo.io,2026:host/unsupported-operation",
+      title: "Document mutations are unavailable on this connection",
+      status: 501,
+    },
+  });
+  return { edit: () => unsupported(), save: unsupported };
+}
+
+/** Choose the backend for the editor's notebook. */
+export function createEditorBackend(
+  options: EditorBackendOptions,
+): NotebookBackend {
+  if (options.static) {
+    return { transport: idle, notebook: uneditable };
+  }
+  if (isWasm()) {
+    return wasm();
+  }
+  const url = afterStartup(options);
+  return {
+    transport:
+      options.transportType === "sse"
+        ? sessionSse(url, options.headers)
+        : websocket(url),
+    notebook: (id) =>
+      new NotebookClient(options.notebookUrl(id), {
+        fetch: withHeaders(options.headers),
+      }),
+    checkHealth: options.checkHealth,
+  };
+}

package/src/core/{websocket → session}/types.ts

@@ -1,12 +1,12 @@
 /* Copyright 2026 Marimo. All rights reserved. */
 
-import type { ConnectionStatus as HostConnectionStatus } from "../host/connection";
+import type { ConnectionStatus as HostConnectionStatus } from "../notebook/replica";
 import type { NotificationMessageData } from "../kernel/messages";
 
 export {
   ConnectionState,
   ConnectionClosedReason as WebSocketClosedReason,
-} from "../host/connection";
+} from "../notebook/replica";
 
 export type ConnectionPhase =
   | NotificationMessageData<"startup-progress">["phase"]

package/src/core/session/useNotebookReplica.ts

@@ -0,0 +1,72 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+/** The editor session's replica, following its backend while mounted. */
+
+import { useStore } from "jotai";
+import { useEffect, useMemo } from "react";
+import useEvent from "react-use-event-hook";
+import type {
+  NotificationMessage,
+  NotificationPayload,
+} from "../kernel/messages";
+import {
+  createNotebookReplica,
+  type NotebookReplica,
+} from "../notebook/replica";
+import { activeNotebookReplicaAtom } from "./active-notebook";
+import { createEditorBackend, type EditorBackendOptions } from "./backend";
+
+type Options = EditorBackendOptions & {
+  onNotification: (notification: NotificationPayload) => void;
+};
+
+export function useNotebookReplica(options: Options): NotebookReplica {
+  const store = useStore();
+  const { static: isStatic, transportType } = options;
+  const replica = useMemo(() => createNotebookReplica(store), [store]);
+  const url = useEvent(options.url);
+  const headers = useEvent(options.headers);
+  const notebookUrl = useEvent(options.notebookUrl);
+  const waitUntilReady = useEvent(options.waitUntilReady);
+  const checkHealth = useEvent(options.checkHealth);
+  const onNotification = useEvent(options.onNotification);
+  // Opens nothing until the replica follows it.
+  const backend = useMemo(
+    () =>
+      createEditorBackend({
+        static: isStatic,
+        transportType,
+        url,
+        headers,
+        notebookUrl,
+        waitUntilReady,
+        checkHealth,
+      }),
+    [
+      isStatic,
+      transportType,
+      url,
+      headers,
+      notebookUrl,
+      waitUntilReady,
+      checkHealth,
+    ],
+  );
+
+  useEffect(() => {
+    store.set(activeNotebookReplicaAtom, replica);
+    const disconnect = replica.connect(backend, {
+      // The host validates the envelope, not the native notification contents.
+      onNotification: (notification) =>
+        onNotification({ data: notification as NotificationMessage }),
+    });
+    return () => {
+      disconnect();
+      if (store.get(activeNotebookReplicaAtom) === replica) {
+        store.set(activeNotebookReplicaAtom, null);
+      }
+    };
+  }, [replica, backend, store, onNotification]);
+
+  return replica;
+}

package/src/core/{websocket → session}/useNotebookSession.tsx

@@ -2,18 +2,15 @@
 
 import { invalidatePackageData } from "@/core/packages/package-data";
 import { useAtom, useAtomValue, useSetAtom, useStore } from "jotai";
-import { useLayoutEffect } from "react";
+import { useEffect, useLayoutEffect, useMemo } from "react";
 import { useErrorBoundary } from "react-error-boundary";
 import { toast } from "@/components/ui/use-toast";
-import { attachmentsAtom } from "@/core/attachments/state";
-import { getNotebook, useCellActions } from "@/core/cells/cells";
-import { applyTransactionChanges } from "@/core/cells/document-changes";
+import { getNotebook, notebookAtom, useCellActions } from "@/core/cells/cells";
+import { followDocument } from "@/core/cells/document";
 import { AUTOCOMPLETER } from "@/core/codemirror/completion/Autocompleter";
-import type {
-  NotificationMessageData,
-  NotificationPayload,
-} from "@/core/kernel/messages";
-import { useNotebookClient } from "@/core/websocket/useNotebookClient";
+import { createDocumentModel } from "@/core/notebook/document";
+import type { NotificationPayload } from "@/core/kernel/messages";
+import { useNotebookReplica } from "@/core/session/useNotebookReplica";
 import { renderHTML } from "@/plugins/core/RenderHTML";
 import {
   handleWidgetMessage,
@@ -26,14 +23,13 @@ import { Logger } from "@/utils/Logger";
 import { reloadSafe } from "@/utils/reload-safe";
 import { useAlertActions } from "../alerts/state";
 import { cacheInfoAtom } from "../cache/requests";
-import { SCRATCH_CELL_ID } from "../cells/ids";
 import { useRunsActions } from "../cells/runs";
 import { focusAndScrollCellOutputIntoView } from "../cells/scrollCellIntoView";
 import {
   resyncBreakpoints,
   setActiveLine,
 } from "../codemirror/cells/debugger-state";
-import type { CellData } from "../cells/types";
+import { notebookCells } from "../cells/utils";
 import { capabilitiesAtom } from "../config/capabilities";
 import { connectionTransportTypeAtom, useSetAppConfig } from "../config/config";
 import { useDataSourceActions } from "../datasets/data-source-connections";
@@ -66,7 +62,7 @@ import { type LayoutState, useLayoutActions } from "../layout/state";
 import { kioskModeAtom } from "../mode";
 import { connectionAtom, startupProgressAtom } from "../network/connection";
 import type { RequestId } from "../network/DeferredRequestRegistry";
-import { runtimeConfigAtom, useRuntimeManager } from "../runtime/config";
+import { useRuntimeManager } from "../runtime/config";
 import { SECRETS_REGISTRY } from "../secrets/request-registry";
 import { isStaticNotebook } from "../static/static-state";
 import {
@@ -76,43 +72,26 @@ import {
 import { useStorageActions } from "../storage/state";
 import { useVariablesActions } from "../variables/state";
 import type { VariableName } from "../variables/types";
-import { isWasm } from "../wasm/utils";
+import { activeDocumentModelAtom } from "./active-notebook";
 import { WebSocketClosedReason, ConnectionState } from "./types";
 
-const SUPPORTS_LAZY_KERNELS = true;
-
-function getExistingCells(): CellData[] | undefined {
-  if (!SUPPORTS_LAZY_KERNELS) {
-    return undefined;
-  }
-
-  // Remove scratch pad
-  return Object.values(getNotebook().cellData).filter(
-    (cell) => cell.id !== SCRATCH_CELL_ID,
-  );
-}
-
 /** Coordinate session startup and apply notifications to the editor. */
 export function useNotebookSession(opts: {
   sessionId: SessionId;
   autoInstantiate: boolean;
-  setCells: (cells: CellData[], layout: LayoutState) => void;
+  /** Called with the layout the kernel reported, after its cells are ready. */
+  onKernelReady?: (layout: LayoutState) => void;
 }) {
-  const { autoInstantiate, sessionId, setCells } = opts;
+  const { autoInstantiate, sessionId, onKernelReady } = opts;
   const { showBoundary } = useErrorBoundary();
+  const store = useStore();
 
-  const { handleCellMessage } = useCellActions();
-  const actionsWithoutMiddleware = useCellActions({ skipMiddleware: true });
-
-  const handleDocumentTransaction = (
-    transaction: NotificationMessageData<"notebook-document-transaction">["transaction"],
-  ) => {
-    applyTransactionChanges(
-      transaction.changes,
-      actionsWithoutMiddleware,
-      () => getNotebook().cellIds.inOrderIds,
-    );
-  };
+  const {
+    handleCellMessage,
+    reconcileDocument,
+    setLastRun,
+    resetCellRuntimes,
+  } = useCellActions();
   const { addCellNotification } = useRunsActions();
   const setKernelState = useSetAtom(kernelStateAtom);
   const setInitialRunCompleted = useSetAtom(initialRunCompletedAtom);
@@ -136,11 +115,9 @@ export function useNotebookSession(opts: {
   const setStableSessionId = useSetAtom(stableSessionIdAtom);
   const setCapabilities = useSetAtom(capabilitiesAtom);
   const runtimeManager = useRuntimeManager();
-  const { notebookId } = useAtomValue(runtimeConfigAtom);
   const transportType = useAtomValue(connectionTransportTypeAtom);
   const setCacheInfo = useSetAtom(cacheInfoAtom);
   const setKernelStartupError = useSetAtom(kernelStartupErrorAtom);
-  const setAttachments = useSetAtom(attachmentsAtom);
   const {
     setNamespaces: setStorageNamespaces,
     filterFromVariables: filterStorageFromVariables,
@@ -148,9 +125,6 @@ export function useNotebookSession(opts: {
 
   const handleNotification = (msg: NotificationPayload) => {
     switch (msg.data.op) {
-      case "attachments":
-        setAttachments(msg.data.attachments);
-        return;
       case "reload":
         reloadSafe();
         return;
@@ -172,20 +146,19 @@ export function useNotebookSession(opts: {
         setInitialRunCompleted(
           Boolean(msg.data.resumed || msg.data.auto_instantiated),
         );
-        const existingCells = getExistingCells();
-
-        handleKernelReady(msg.data, {
+        const layout = handleKernelReady(msg.data, {
           autoInstantiate,
-          setCells,
+          cells: notebookCells(store.get(notebookAtom)),
+          setLastRun,
+          resetCellRuntimes,
           setLayoutData,
           setAppConfig,
           setCapabilities,
           setKernelState,
           onError: showBoundary,
-          existingCells,
         });
+        onKernelReady?.(layout);
         if (!msg.data.resumed) {
-          setAttachments([]);
           // A freshly started kernel may expose a different environment
           // (e.g. new env vars); re-run discovery instead of reusing stale
           // suggestions from the previous kernel instance.
@@ -271,13 +244,18 @@ export function useNotebookSession(opts: {
         filterStorageFromVariables(variableNames);
         return;
       case "variable-values":
-        setMetadata(
-          msg.data.variables.map((v) => ({
+        setMetadata([
+          ...(msg.data.removed ?? []).map((v) => ({
+            name: v.name as VariableName,
+            dataType: undefined,
+            value: undefined,
+          })),
+          ...msg.data.variables.map((v) => ({
             name: v.name as VariableName,
             dataType: v.datatype,
             value: v.value,
           })),
-        );
+        ]);
         return;
       case "alert":
         toast({
@@ -398,7 +376,7 @@ export function useNotebookSession(opts: {
         );
         return;
       case "notebook-document-transaction":
-        handleDocumentTransaction(msg.data.transaction);
+        // Edits reach the notebook through the host document.
         return;
       case "consumer-capabilities":
         setKioskMode(!msg.data.consumer_capabilities.edit);
@@ -408,7 +386,6 @@ export function useNotebookSession(opts: {
     }
   };
 
-  const store = useStore();
   const onMessageError = (error: unknown) => {
     toast({
       title: "Failed to handle message",
@@ -416,25 +393,23 @@ export function useNotebookSession(opts: {
       variant: "danger",
     });
   };
-  const notebook = useNotebookClient({
-    sessionId,
+  const replica = useNotebookReplica({
     static: isStaticNotebook(),
     transportType,
-    messageFormat:
-      notebookId && transportType === "websocket" ? "host" : "session",
     url: () =>
       transportType === "sse"
         ? runtimeManager.getSseURL(sessionId).toString()
         : runtimeManager.getWsURL(sessionId).toString(),
     // SSE uses headers for auth. WebSockets carry the token in the URL.
     headers: () => runtimeManager.headers(),
+    notebookUrl: (notebookId) => runtimeManager.getNotebookURL(notebookId),
+    // Only the host's transports read the URL, so static and WASM never wait.
     waitUntilReady: async () => {
-      if (!isStaticNotebook() && !isWasm() && !runtimeManager.isSameOrigin) {
+      if (!runtimeManager.isSameOrigin) {
         await runtimeManager.waitForHealthy();
       }
     },
     checkHealth: () => runtimeManager.reconcileFromHealth(),
-    onMessageError,
     onNotification: (notification) => {
       try {
         handleNotification(notification);
@@ -445,14 +420,42 @@ export function useNotebookSession(opts: {
     },
   });
 
+  const documentModel = useMemo(
+    () =>
+      createDocumentModel({
+        store,
+        documentAtom: replica.documentAtom,
+        backendAtom: replica.backendAtom,
+      }),
+    [store, replica],
+  );
+  // Edits outlive reconnects and end with the session. The notebook's cells
+  // follow the document for as long as the session is mounted.
+  useEffect(() => {
+    const controller = new AbortController();
+    documentModel.start({ signal: controller.signal });
+    store.set(activeDocumentModelAtom, documentModel);
+    followDocument({
+      store,
+      model: documentModel,
+      reconcile: reconcileDocument,
+      signal: controller.signal,
+    });
+    return () => {
+      controller.abort();
+      if (store.get(activeDocumentModelAtom) === documentModel) {
+        store.set(activeDocumentModelAtom, null);
+      }
+    };
+  }, [documentModel, store, reconcileDocument]);
+
   // Subscribe before passive connection setup: static/WASM transports can emit synchronously.
   useLayoutEffect(() => {
     const updateConnection = () => {
-      const status = store.get(notebook.connectionAtom);
+      const status = store.get(replica.connectionAtom);
       if (status.state === ConnectionState.NOT_STARTED) {
         return;
       }
-      setAttachments([]);
       if (status.state === ConnectionState.OPEN) {
         updateStartupProgress(null);
         return; // kernel-ready, not the socket, establishes session readiness.
@@ -480,17 +483,16 @@ export function useNotebookSession(opts: {
         return status;
       });
     };
-    const unsubscribe = store.sub(notebook.connectionAtom, updateConnection);
+    const unsubscribe = store.sub(replica.connectionAtom, updateConnection);
     updateConnection();
     return unsubscribe;
   }, [
-    notebook,
+    replica,
     store,
     setConnection,
-    setAttachments,
     setKernelStartupError,
     updateStartupProgress,
   ]);
 
-  return { notebook, connection, reconnect: notebook.reconnect };
+  return { replica, documentModel, connection, reconnect: replica.reconnect };
 }

package/src/core/state/__tests__/jotai.test.tsx

@@ -0,0 +1,108 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+import { act, renderHook } from "@testing-library/react";
+import { atom, Provider } from "jotai";
+import type React from "react";
+import { describe, expect, it, vi } from "vitest";
+import { store, useJotaiEffect, waitFor } from "../jotai";
+
+const wrapper = ({ children }: { children: React.ReactNode }) => (
+  <Provider store={store}>{children}</Provider>
+);
+
+describe("waitFor function", () => {
+  it("should resolve when the atom satisfies the predicate", async () => {
+    const testAtom = atom(0);
+    store.set(testAtom, 10);
+
+    const result = await waitFor(testAtom, (value) => value === 10);
+    expect(result).toBe(10);
+  });
+
+  it("should resolve when the atom changes to satisfy the predicate", async () => {
+    const testAtom = atom(0);
+
+    setTimeout(() => store.set(testAtom, 15), 5);
+    setTimeout(() => store.set(testAtom, 20), 10);
+
+    const result = await waitFor(testAtom, (value) => value === 20);
+    expect(result).toBe(20);
+  });
+});
+
+describe("useJotaiEffect", () => {
+  it("calls the effect with the new and previous values", () => {
+    const testAtom = atom(0);
+    const effect = vi.fn();
+    renderHook(() => useJotaiEffect(testAtom, effect), { wrapper });
+
+    act(() => store.set(testAtom, 1));
+    act(() => store.set(testAtom, 2));
+
+    expect(effect.mock.calls).toEqual([
+      [1, 0],
+      [2, 1],
+    ]);
+  });
+
+  it("calls only the latest callback after a rerender", () => {
+    const testAtom = atom(0);
+    const first = vi.fn();
+    const second = vi.fn();
+    const { rerender } = renderHook(
+      ({ effect }) => useJotaiEffect(testAtom, effect),
+      { initialProps: { effect: first }, wrapper },
+    );
+
+    rerender({ effect: second });
+    act(() => store.set(testAtom, 1));
+
+    expect(first).not.toHaveBeenCalled();
+    expect(second).toHaveBeenCalledTimes(1);
+  });
+
+  it("subscribes once per rerender with an inline callback", () => {
+    const testAtom = atom(0);
+    const calls: number[] = [];
+    const { rerender } = renderHook(
+      () => useJotaiEffect(testAtom, (value) => calls.push(value)),
+      { wrapper },
+    );
+
+    rerender();
+    rerender();
+    act(() => store.set(testAtom, 1));
+
+    expect(calls).toEqual([1]);
+  });
+
+  it("stops listening to the previous atom when the atom changes", () => {
+    const firstAtom = atom(0);
+    const secondAtom = atom(0);
+    const effect = vi.fn();
+    const { rerender } = renderHook(
+      ({ target }) => useJotaiEffect(target, effect),
+      { initialProps: { target: firstAtom }, wrapper },
+    );
+
+    rerender({ target: secondAtom });
+    act(() => store.set(firstAtom, 1));
+    expect(effect).not.toHaveBeenCalled();
+
+    act(() => store.set(secondAtom, 1));
+    expect(effect).toHaveBeenCalledTimes(1);
+  });
+
+  it("unsubscribes on unmount", () => {
+    const testAtom = atom(0);
+    const effect = vi.fn();
+    const { unmount } = renderHook(() => useJotaiEffect(testAtom, effect), {
+      wrapper,
+    });
+
+    unmount();
+    act(() => store.set(testAtom, 1));
+
+    expect(effect).not.toHaveBeenCalled();
+  });
+});

package/src/core/state/jotai.ts

@@ -42,7 +42,7 @@ export function useJotaiEffect<T>(
   const store = useStore();
   useEffect(() => {
     let prevValue = store.get(atom);
-    store.sub(atom, () => {
+    return store.sub(atom, () => {
       const value = store.get(atom);
       effect(value, prevValue);
       prevValue = value;

package/src/core/static/__tests__/files.test.ts

@@ -654,19 +654,47 @@ describe("error handling and edge cases", () => {
     unpatch();
   });
 
-  it("should handle malformed URLs gracefully", async () => {
-    const virtualFiles = {};
-    const unpatch = patchFetch(virtualFiles);
+  it("forwards relative URLs and options to the original fetch", async () => {
+    const response = new Response("relative content");
+    const originalFetch = vi.spyOn(window, "fetch").mockResolvedValue(response);
+    const unpatch = patchFetch({});
+    const options = { headers: { Accept: "text/plain" } };
+
+    try {
+      expect(await window.fetch("relative-file.txt", options)).toBe(response);
+      expect(originalFetch).toHaveBeenCalledExactlyOnceWith(
+        "relative-file.txt",
+        options,
+      );
+    } finally {
+      unpatch();
+      originalFetch.mockRestore();
+    }
+  });
 
-    // Mock Logger.error to avoid test output
+  it("preserves the original fetch rejection for malformed URLs", async () => {
+    const error = new TypeError("Invalid URL");
+    const originalFetch = vi.spyOn(window, "fetch").mockRejectedValue(error);
     const loggerSpy = vi
       .spyOn(console, "error")
       .mockImplementation(Functions.NOOP);
-
-    // This should catch the error and fallback to original fetch
-    await window.fetch("not-a-valid-url").catch(() => undefined);
-
-    unpatch();
-    loggerSpy.mockRestore();
+    const unpatch = patchFetch({});
+    const options = { headers: { Accept: "text/plain" } };
+
+    try {
+      await expect(window.fetch("http://[", options)).rejects.toBe(error);
+      expect(originalFetch).toHaveBeenCalledExactlyOnceWith(
+        "http://[",
+        options,
+      );
+      expect(loggerSpy).toHaveBeenCalledWith(
+        "Error parsing URL",
+        expect.any(TypeError),
+      );
+    } finally {
+      unpatch();
+      originalFetch.mockRestore();
+      loggerSpy.mockRestore();
+    }
   });
 });

package/src/core/static/files.ts

@@ -125,6 +125,8 @@ function withoutLeadingDot(path: string): string {
 /**
  * Resolve a URL to a blob URL if it's a virtual file, for use with dynamic import().
  * Unlike fetch, import() can't be patched, so we need to convert data URLs to blob URLs.
+ * The caller owns a returned URL that differs from the input and must revoke it
+ * once no module consumers or their cleanup callbacks need it.
  *
  * @returns The original URL if not a virtual file, or a blob URL if it is
  */