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

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

@@ -1,19 +1,20 @@
 /* Copyright 2026 Marimo. All rights reserved. */
 
 import { invalidatePackageData } from "@/core/packages/package-data";
-import { useAtom, useAtomValue, useSetAtom } from "jotai";
-import { useRef } from "react";
+import { useAtom, useAtomValue, useSetAtom, useStore } from "jotai";
+import { useEffect, useLayoutEffect, useMemo } from "react";
 import { useErrorBoundary } from "react-error-boundary";
 import { toast } from "@/components/ui/use-toast";
-import { getNotebook, useCellActions } from "@/core/cells/cells";
-import { applyTransactionChanges } from "@/core/cells/document-changes";
+import {
+  getNotebook,
+  notebookAtom,
+  useReducerActions,
+} 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 { TRANSPORT_EXHAUSTED_REASON } from "@/core/websocket/transports/ws";
-import { useConnectionTransport } from "@/core/websocket/useWebSocket";
+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,
@@ -21,25 +22,20 @@ import {
 } from "@/plugins/impl/anywidget/registry";
 import { logNever } from "@/utils/assertNever";
 import { prettyError } from "@/utils/errors";
-import {
-  type JsonString,
-  safeExtractSetUIElementMessageBuffers,
-} from "@/utils/json/base64";
-import { jsonParseWithSpecialChar } from "@/utils/json/json-parser";
+import { safeExtractSetUIElementMessageBuffers } from "@/utils/json/base64";
 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 { connectionTransportTypeAtom } from "../config/config";
 import { useDataSourceActions } from "../datasets/data-source-connections";
 import { invalidateDataSourceDiscovery } from "../datasets/data-source-discovery";
 import type { ConnectionName } from "../datasets/engines";
@@ -66,9 +62,10 @@ import {
   stableSessionIdAtom,
 } from "../kernel/session";
 import { initialRunCompletedAtom, kernelStateAtom } from "../kernel/state";
-import { type LayoutState, useLayoutActions } from "../layout/state";
+import { useLayoutActions } from "../layout/state";
 import { kioskModeAtom } from "../mode";
 import { connectionAtom, startupProgressAtom } from "../network/connection";
+import { savedLayoutAtom } from "../saving/state";
 import type { RequestId } from "../network/DeferredRequestRegistry";
 import { useRuntimeManager } from "../runtime/config";
 import { SECRETS_REGISTRY } from "../secrets/request-registry";
@@ -80,145 +77,28 @@ import {
 import { useStorageActions } from "../storage/state";
 import { useVariablesActions } from "../variables/state";
 import type { VariableName } from "../variables/types";
-import { isWasm } from "../wasm/utils";
-import {
-  type ConnectionStatus,
-  WebSocketClosedReason,
-  ConnectionState,
-} from "./types";
-
-const SUPPORTS_LAZY_KERNELS = true;
-
-// All MARIMO_* reasons except TRANSPORT_EXHAUSTED are emitted by the backend
-// (marimo/_server/api/endpoints/ws_endpoint.py and ws/*.py). Keep in sync with
-// the backend literals. Over the SSE transport these arrive as in-band
-// `close` events rather than WebSocket close frames, so every terminal
-// reason must be classified here — an unclassified reason falls into the
-// retry path.
-export type CloseReason =
-  | "MARIMO_NO_FILE_KEY"
-  | "MARIMO_NO_SESSION_ID"
-  | "MARIMO_NO_SESSION"
-  | "MARIMO_SHUTDOWN"
-  | "MARIMO_KERNEL_STARTUP_ERROR"
-  | "MARIMO_UNAUTHORIZED"
-  | "MARIMO_KIOSK_NOT_ALLOWED"
-  | typeof TRANSPORT_EXHAUSTED_REASON;
-
-export type CloseDecision =
-  | { kind: "terminal"; status: ConnectionStatus; closeTransport: boolean }
-  | { kind: "gave-up"; status: ConnectionStatus }
-  | { kind: "retry"; status: ConnectionStatus };
-
-export function classifyCloseEvent(event: { reason?: string }): CloseDecision {
-  switch (event.reason as CloseReason | undefined) {
-    case TRANSPORT_EXHAUSTED_REASON:
-      return {
-        kind: "gave-up",
-        status: {
-          state: ConnectionState.CLOSED,
-          code: WebSocketClosedReason.KERNEL_DISCONNECTED,
-          reason: "kernel not found",
-        },
-      };
-    case "MARIMO_NO_FILE_KEY":
-    case "MARIMO_NO_SESSION_ID":
-    case "MARIMO_NO_SESSION":
-    case "MARIMO_SHUTDOWN":
-      return {
-        kind: "terminal",
-        status: {
-          state: ConnectionState.CLOSED,
-          code: WebSocketClosedReason.KERNEL_DISCONNECTED,
-          reason: "kernel not found",
-        },
-        closeTransport: true,
-      };
-    case "MARIMO_UNAUTHORIZED":
-      return {
-        kind: "terminal",
-        status: {
-          state: ConnectionState.CLOSED,
-          code: WebSocketClosedReason.KERNEL_DISCONNECTED,
-          reason: "not authorized",
-        },
-        closeTransport: true,
-      };
-    case "MARIMO_KIOSK_NOT_ALLOWED":
-      return {
-        kind: "terminal",
-        status: {
-          state: ConnectionState.CLOSED,
-          code: WebSocketClosedReason.KERNEL_DISCONNECTED,
-          reason: "kiosk mode is not available for this session",
-        },
-        closeTransport: true,
-      };
-    case "MARIMO_KERNEL_STARTUP_ERROR":
-      return {
-        kind: "terminal",
-        status: {
-          state: ConnectionState.CLOSED,
-          code: WebSocketClosedReason.KERNEL_STARTUP_ERROR,
-          reason: "Failed to start kernel sandbox",
-        },
-        closeTransport: true,
-      };
-    default:
-      // Empty/undefined reasons are normal transient closes. Anything else is
-      // an unknown server reason; warn so a new MARIMO_* reason doesn't fall
-      // silently into the retry path.
-      if (event.reason) {
-        logNever(event.reason as never);
-      }
-  }
-
-  return {
-    kind: "retry",
-    status: { state: ConnectionState.CONNECTING },
-  };
-}
-
-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,
-  );
-}
+import { activeDocumentModelAtom } from "./active-notebook";
+import { embeddedSnapshotAtom } from "./embedded";
+import { WebSocketClosedReason, ConnectionState } from "./types";
 
-/**
- * Creates a connection to the Marimo kernel and handles incoming messages.
- */
-export function useMarimoKernelConnection(opts: {
+/** Coordinate session startup and apply notifications to the editor. */
+export function useNotebookSession(opts: {
   sessionId: SessionId;
   autoInstantiate: boolean;
-  setCells: (cells: CellData[], layout: LayoutState) => void;
 }) {
-  // Track whether we want to try reconnecting.
-  const shouldTryReconnecting = useRef<boolean>(true);
-  const { autoInstantiate, sessionId, setCells } = opts;
+  const { autoInstantiate, sessionId } = 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,
+  } = useReducerActions();
   const { addCellNotification } = useRunsActions();
   const setKernelState = useSetAtom(kernelStateAtom);
   const setInitialRunCompleted = useSetAtom(initialRunCompletedAtom);
-  const setAppConfig = useSetAppConfig();
   const { setVariables, setMetadata } = useVariablesActions();
   const { addColumnPreview } = useDatasetsActions();
   const { addDatasets, filterDatasetsFromVariables } = useDatasetsActions();
@@ -246,8 +126,7 @@ export function useMarimoKernelConnection(opts: {
     filterFromVariables: filterStorageFromVariables,
   } = useStorageActions();
 
-  const handleMessage = (e: MessageEvent<JsonString<NotificationPayload>>) => {
-    const msg = jsonParseWithSpecialChar(e.data);
+  const handleNotification = (msg: NotificationPayload) => {
     switch (msg.data.op) {
       case "reload":
         reloadSafe();
@@ -270,18 +149,17 @@ export function useMarimoKernelConnection(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,
         });
+        store.set(savedLayoutAtom, layout);
         if (!msg.data.resumed) {
           // A freshly started kernel may expose a different environment
           // (e.g. new env vars); re-run discovery instead of reusing stale
@@ -368,13 +246,18 @@ export function useMarimoKernelConnection(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({
@@ -495,7 +378,7 @@ export function useMarimoKernelConnection(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);
@@ -505,107 +388,85 @@ export function useMarimoKernelConnection(opts: {
     }
   };
 
-  const tryReconnecting = (code?: number, reason?: string) => {
-    // If not properly gated, we could try reconnecting forever if the
-    // issue is not transient. So we want to try reconnecting only once after an
-    // open connection is closed.
-    if (shouldTryReconnecting.current) {
-      shouldTryReconnecting.current = false;
-      ws.reconnect(code, reason);
-    }
+  const onMessageError = (error: unknown) => {
+    toast({
+      title: "Failed to handle message",
+      description: prettyError(error),
+      variant: "danger",
+    });
   };
-
-  // Manual reconnect. Probes /health first to fail fast when the runtime
-  // is unreachable, instead of waiting on partysocket's retry budget.
-  const reconnect = async () => {
-    if (
-      ws.readyState === WebSocket.OPEN ||
-      ws.readyState === WebSocket.CONNECTING
-    ) {
-      return;
-    }
-    shouldTryReconnecting.current = true;
-    setKernelStartupError(null);
-    setConnection({ state: ConnectionState.CONNECTING });
-    const healthy = await runtimeManager.reconcileFromHealth();
-    if (!healthy) {
-      shouldTryReconnecting.current = false;
-      setConnection({
-        state: ConnectionState.CLOSED,
-        code: WebSocketClosedReason.KERNEL_DISCONNECTED,
-        reason: "kernel not found",
-      });
-      return;
-    }
-    ws.reconnect();
-  };
-
-  const ws = useConnectionTransport({
+  const replica = useNotebookReplica({
     static: isStaticNotebook(),
+    snapshot: useAtomValue(embeddedSnapshotAtom),
     transportType,
-    /**
-     * Unique URL for this session.
-     */
     url: () =>
       transportType === "sse"
         ? runtimeManager.getSseURL(sessionId).toString()
         : runtimeManager.getWsURL(sessionId).toString(),
-    /**
-     * Auth headers for the SSE transport. WebSockets cannot send headers
-     * and instead put the access token in the URL.
-     */
+    // SSE uses headers for auth. WebSockets carry the token in the URL.
     headers: () => runtimeManager.headers(),
-
-    /**
-     * The transport is open; kernel-ready establishes session readiness.
-     */
-    onOpen: async () => {
-      updateStartupProgress(null);
-      // If we are open, we can reset our reconnecting flag.
-      shouldTryReconnecting.current = true;
-    },
-
-    /**
-     * Wait to connect, in case the remote kernel still starting up.
-     */
-    waitToConnect: async () => {
-      if (isStaticNotebook()) {
-        return;
+    notebookUrl: (notebookId) => runtimeManager.getNotebookURL(notebookId),
+    // Only the host's transports read the URL, so static and WASM never wait.
+    waitUntilReady: async () => {
+      if (!runtimeManager.isSameOrigin) {
+        await runtimeManager.waitForHealthy();
       }
-      if (isWasm()) {
-        return;
-      }
-
-      if (runtimeManager.isSameOrigin) {
-        return;
-      }
-      await runtimeManager.waitForHealthy();
     },
-
-    /**
-     * Handle messages sent by the kernel.
-     */
-    onMessage: (e) => {
+    checkHealth: () => runtimeManager.reconcileFromHealth(),
+    onNotification: (notification) => {
       try {
-        handleMessage(e);
+        handleNotification(notification);
       } catch (error) {
-        Logger.error("Failed to handle message", e.data, error);
-        toast({
-          title: "Failed to handle message",
-          description: prettyError(error),
-          variant: "danger",
-        });
+        Logger.error("Failed to handle notification", notification, error);
+        onMessageError(error);
       }
     },
+  });
+
+  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]);
 
-    /**
-     * Handle a close event. We may want to reconnect.
-     */
-    onClose: (e) => {
-      Logger.warn("WebSocket closed", e.code, e.reason);
-      const decision = classifyCloseEvent(e);
+  // Subscribe before passive connection setup: static/WASM transports can emit synchronously.
+  useLayoutEffect(() => {
+    const updateConnection = () => {
+      const status = store.get(replica.connectionAtom);
+      if (status.state === ConnectionState.NOT_STARTED) {
+        return;
+      }
+      if (status.state === ConnectionState.OPEN) {
+        updateStartupProgress(null);
+        return; // kernel-ready, not the socket, establishes session readiness.
+      }
+      if (status.state === ConnectionState.CONNECTING) {
+        setKernelStartupError(null);
+      }
       setConnection((previous) => {
-        const status = decision.status;
         if (
           status.state === ConnectionState.CLOSED &&
           status.code === WebSocketClosedReason.KERNEL_STARTUP_ERROR &&
@@ -624,31 +485,17 @@ export function useMarimoKernelConnection(opts: {
         }
         return status;
       });
-      if (decision.kind === "terminal" && decision.closeTransport) {
-        ws.close(); // close to prevent reconnecting
-        return;
-      }
-      if (decision.kind === "retry") {
-        // Session should be valid
-        // - browser tab might have been closed or re-opened
-        // - computer might have just woken from sleep
-        tryReconnecting(e.code, e.reason);
-      }
-    },
-
-    /**
-     * When we encounter an error, we should close the connection.
-     */
-    onError: (e) => {
-      Logger.warn("WebSocket error", e);
-      setConnection({
-        state: ConnectionState.CLOSED,
-        code: WebSocketClosedReason.KERNEL_DISCONNECTED,
-        reason: "kernel not found",
-      });
-      tryReconnecting();
-    },
-  });
-
-  return { connection, reconnect };
+    };
+    const unsubscribe = store.sub(replica.connectionAtom, updateConnection);
+    updateConnection();
+    return unsubscribe;
+  }, [
+    replica,
+    store,
+    setConnection,
+    setKernelStartupError,
+    updateStartupProgress,
+  ]);
+
+  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();
+    }
   });
 });