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

package/src/core/edit-app.tsx

@@ -5,8 +5,8 @@ import { Tooltip } from "radix-ui";
 
 const TooltipProvider = Tooltip.Provider;
 
-import { useAtomValue, useSetAtom } from "jotai";
-import { useEffect } from "react";
+import { useAtomValue } from "jotai";
+import { lazy, useEffect } from "react";
 import { NotStartedConnectionAlert } from "@/components/editor/alerts/connecting-alert";
 import { Controls } from "@/components/editor/controls/Controls";
 import { SandboxController } from "@/components/editor/chrome/panels/sandbox-controller";
@@ -27,23 +27,22 @@ import { useSetOutputsExpanded } from "../components/editor/actions/useSetOutput
 import { CellArray } from "../components/editor/renderers/cell-array";
 import { CellsRenderer } from "../components/editor/renderers/cells-renderer";
 import { useHotkey } from "../hooks/useHotkey";
-import {
-  hasCellsAtom,
-  notebookIsRunningAtom,
-  numColumnsAtom,
-  useCellActions,
-} from "./cells/cells";
+import { hasCellsAtom, notebookIsRunningAtom } from "./cells/cells";
+import { numColumnsAtom } from "./cells/layout";
+import { useCellActions } from "./cells/actions";
 import type { AppConfig, UserConfig } from "./config/config-schema";
 import { RuntimeState } from "./kernel/RuntimeState";
-import { getSessionId } from "./kernel/session";
 import { useTogglePresenting } from "./layout/useTogglePresenting";
 import { viewStateAtom } from "./mode";
 import { useRequestClient } from "./network/requests";
 import { useFilename } from "./saving/filename";
-import { lastSavedNotebookAtom } from "./saving/state";
-import { useMarimoKernelConnection } from "./websocket/useMarimoKernelConnection";
+import type { ConnectionStatus } from "./session/types";
+
+const LazyCommandPalette = lazy(
+  () => import("../components/editor/controls/command-palette"),
+);
 
-interface AppProps {
+export interface EditAppProps {
   /**
    * The user config.
    */
@@ -56,19 +55,24 @@ interface AppProps {
    * If true, the floating controls will be hidden.
    */
   hideControls?: boolean;
+  /** The session's connection, which the page owns with the document model. */
+  connection: ConnectionStatus;
+  reconnect: () => Promise<void>;
 }
 
-export const EditApp: React.FC<AppProps> = ({
+/** The notebook editor. Rendered inside the page's `DocumentProvider`. */
+export const EditApp: React.FC<EditAppProps> = ({
   userConfig,
   appConfig,
   hideControls = false,
+  connection,
+  reconnect,
 }) => {
-  const { setCells, mergeAllColumns, collapseAllCells, expandAllCells } =
+  const { mergeAllColumns, collapseAllCells, expandAllCells } =
     useCellActions();
   const viewState = useAtomValue(viewStateAtom);
   const numColumns = useAtomValue(numColumnsAtom);
   const filename = useFilename();
-  const setLastSavedNotebook = useSetAtom(lastSavedNotebookAtom);
   const { sendComponentValues, sendInterrupt } = useRequestClient();
 
   const isEditing = viewState.mode === "edit";
@@ -81,19 +85,7 @@ export const EditApp: React.FC<AppProps> = ({
     return () => {
       RuntimeState.INSTANCE.stop();
     };
-  }, []);
-
-  const { connection, reconnect } = useMarimoKernelConnection({
-    autoInstantiate: userConfig.runtime.auto_instantiate,
-    setCells: (cells, layout) => {
-      setCells(cells);
-      const names = cells.map((cell) => cell.name);
-      const codes = cells.map((cell) => cell.code);
-      const configs = cells.map((cell) => cell.config);
-      setLastSavedNotebook({ names, codes, configs, layout });
-    },
-    sessionId: getSessionId(),
-  });
+  }, [sendComponentValues]);
 
   // Update document title whenever filename or app_title changes
   useEffect(() => {
@@ -228,6 +220,7 @@ export const EditApp: React.FC<AppProps> = ({
           />
         </TooltipProvider>
       )}
+      <LazyCommandPalette />
     </>
   );
 };

package/src/core/errors/error-entries.ts

@@ -80,7 +80,7 @@ export function getCellErrorEntries(store: JotaiStore): CellErrorEntry[] {
   const { cellIds, cellRuntime, cellData } = store.get(notebookAtom);
   const entries: CellErrorEntry[] = [];
 
-  for (const [cellIndex, cellId] of cellIds.inOrderIds.entries()) {
+  for (const [cellIndex, cellId] of cellIds.entries()) {
     const cell = cellRuntime[cellId];
     const output = cell.output;
     if (!output || !isErrorMime(output.mimetype)) {
@@ -122,6 +122,9 @@ export function describeError(error: MarimoError): string {
   if (error.type === "import-star") {
     return error.msg;
   }
+  if (error.type === "lazy-import") {
+    return error.msg;
+  }
   if (error.type === "ancestor-stopped") {
     return error.msg;
   }

package/src/core/export/hooks.ts

@@ -16,7 +16,7 @@ import { connectionAtom } from "../network/connection";
 import { useRequestClient } from "../network/requests";
 import type { UpdateCellOutputsRequest } from "../network/types";
 import { VirtualFileTracker } from "../static/virtual-file-tracker";
-import { ConnectionState } from "../websocket/types";
+import { ConnectionState } from "../session/types";
 
 const DELAY = 5000; // 5 seconds;
 

package/src/core/export/layout.ts

@@ -1,9 +1,9 @@
 /* Copyright 2026 Marimo. All rights reserved. */
 
-import { flushDocumentChanges } from "@/core/cells/document-changes";
+import { flushDocument } from "@/core/cells/document";
 import { getSerializedLayout } from "@/core/layout/layout";
 
 export async function getExportLayout() {
-  await flushDocumentChanges();
+  await flushDocument();
   return getSerializedLayout();
 }

package/src/core/hotkeys/actions.ts

@@ -2,6 +2,7 @@
 import { atom, useAtomValue, useSetAtom } from "jotai";
 import useEvent from "react-use-event-hook";
 import type { HotkeyAction } from "@/core/hotkeys/hotkeys";
+import { Objects } from "@/utils/objects";
 
 /**
  * Map of registered keyboard shortcuts and their callbacks.
@@ -29,10 +30,7 @@ export function useSetRegisteredAction() {
       },
     ),
     unregisterAction: useEvent((shortcut: HotkeyAction) => {
-      set((actions) => {
-        const { [shortcut]: unused, ...rest } = actions;
-        return rest;
-      });
+      set((actions) => Objects.omit(actions, [shortcut]));
     }),
   };
 }

package/src/core/islands/__tests__/bootstrap.test.ts

@@ -1,22 +1,29 @@
 /* Copyright 2026 Marimo. All rights reserved. */
 
+import type { NotebookTransport } from "@marimo-team/client";
 import { createStore } from "jotai";
 import { afterEach, describe, expect, it, vi } from "vitest";
+import { cellId } from "@/__tests__/branded";
+import { notebookAtom } from "@/core/cells/cells";
+import {
+  fakeBackend,
+  fixtureMessages,
+} from "@/core/notebook/replica/__tests__/fixtures";
 import { RuntimeState } from "@/core/kernel/RuntimeState";
 import { initializeIslands } from "../bootstrap";
 import type { IslandsPyodideBridge } from "../bridge";
 import { islandsHydratedAtom, islandsPendingInitialRunsAtom } from "../state";
-import type { IslandsKernelMessage } from "../worker/worker";
-
-function completedRun(sessionGeneration: number): IslandsKernelMessage {
-  return {
-    sessionGeneration,
-    message: JSON.stringify({
-      op: "completed-run",
-      data: { op: "completed-run", run_id: null },
-    }) as IslandsKernelMessage["message"],
-  };
-}
+import type { NotificationMessage } from "@/core/kernel/messages";
+
+type Consumer = (
+  notification: NotificationMessage,
+  sessionGeneration: number,
+) => void;
+
+const completedRun = {
+  op: "completed-run",
+  run_id: null,
+} as NotificationMessage;
 
 describe("initializeIslands", () => {
   afterEach(vi.restoreAllMocks);
@@ -26,12 +33,13 @@ describe("initializeIslands", () => {
     root.innerHTML =
       '<marimo-island data-app-id="app-1" data-reactive="true"></marimo-island>';
 
-    let consumeMessage: ((message: IslandsKernelMessage) => void) | undefined;
+    let consumeMessage: Consumer | undefined;
     const bridge = {
-      consumeMessages: (consumer: (message: IslandsKernelMessage) => void) => {
+      consumeNotifications: (consumer: Consumer) => {
         consumeMessage = consumer;
       },
       sendComponentValues: vi.fn(),
+      notebookBackend: fakeBackend(),
     } as unknown as IslandsPyodideBridge;
     const store = createStore();
     vi.spyOn(RuntimeState.INSTANCE, "start").mockImplementation(
@@ -48,14 +56,58 @@ describe("initializeIslands", () => {
     expect(store.get(islandsHydratedAtom)).toBe(false);
     store.set(islandsPendingInitialRunsAtom, new Set([1, 2]));
 
-    consumeMessage?.(completedRun(1));
+    consumeMessage?.(completedRun, 1);
     expect(store.get(islandsHydratedAtom)).toBe(false);
 
-    consumeMessage?.(completedRun(1));
-    consumeMessage?.(completedRun(0));
+    consumeMessage?.(completedRun, 1);
+    consumeMessage?.(completedRun, 0);
     expect(store.get(islandsHydratedAtom)).toBe(false);
 
-    consumeMessage?.(completedRun(2));
+    consumeMessage?.(completedRun, 2);
     expect(store.get(islandsHydratedAtom)).toBe(true);
   });
 });
+
+describe("initializeIslands document", () => {
+  afterEach(vi.restoreAllMocks);
+
+  it("fills the notebook's cells from the managed session's host document", async () => {
+    const root = document.createElement("div");
+    root.innerHTML =
+      '<marimo-island data-app-id="app-1" data-reactive="true"></marimo-island>';
+    let handlers: Parameters<NotebookTransport>[1] | undefined;
+    const transport: NotebookTransport = (_attempt, opened) =>
+      new Promise(() => {
+        handlers = opened;
+        opened.open();
+      });
+    const bridge = {
+      consumeNotifications: vi.fn(),
+      sendComponentValues: vi.fn(),
+      notebookBackend: fakeBackend({ transport }),
+    } as unknown as IslandsPyodideBridge;
+    const store = createStore();
+    vi.spyOn(RuntimeState.INSTANCE, "start").mockImplementation(
+      () => undefined,
+    );
+
+    await initializeIslands({
+      bridge,
+      store,
+      root,
+      autoInitializePlugins: false,
+    });
+    await vi.waitFor(() => expect(handlers).toBeDefined());
+    for (const message of fixtureMessages("notebook-events.sse")) {
+      handlers?.message(message);
+    }
+
+    const notebook = store.get(notebookAtom);
+    expect(notebook.cellIds).toEqual(["c1", "c3", "k7f2"]);
+    expect(notebook.cellData[cellId("c3")].code).toBe("x = 21");
+    expect(notebook.cellRuntime[cellId("c3")]).toBeDefined();
+
+    handlers?.message({ event: "reset", data: {} });
+    expect(store.get(notebookAtom).cellIds).toEqual(["c1", "c3", "k7f2"]);
+  });
+});

package/src/core/islands/__tests__/bridge.test.ts

@@ -10,7 +10,6 @@ import {
 } from "@/__tests__/branded";
 import { notebookAtom } from "@/core/cells/cells";
 import { islandsPendingInitialRunsAtom } from "@/core/islands/state";
-import type { IslandsKernelMessage } from "../worker/worker";
 
 type Base64String = components["schemas"]["Base64String"];
 interface TestIslandApp {
@@ -89,9 +88,6 @@ vi.mock("@/core/wasm/rpc", () => ({
         startSession: mockStartSessionRequest,
         stopSession: mockStopSessionRequest,
       },
-      send: {
-        consumerReady: vi.fn(),
-      },
     },
     addMessageListener: vi.fn(
       (name: string, listener: (payload: never) => void) => {
@@ -177,17 +173,22 @@ describe("IslandsPyodideBridge", () => {
     });
   }
 
-  function sendKernelMessage(sessionGeneration: number) {
-    const listener = mockMessageListeners.get("kernelMessage");
+  function sendNotebookEvent(sessionGeneration: number, frame: unknown) {
+    const listener = mockMessageListeners.get("notebookEvent");
     if (!listener) {
-      throw new Error("Missing kernel message listener");
+      throw new Error("Missing notebook event listener");
     }
-    const message = {
-      message: "" as IslandsKernelMessage["message"],
-      sessionGeneration,
-    };
-    listener(message as never);
-    return message;
+    listener({ event: JSON.stringify(frame), sessionGeneration } as never);
+  }
+
+  /** Send a notification and return the call the consumer should receive. */
+  function sendKernelMessage(sessionGeneration: number) {
+    const notification = { op: "interrupted" };
+    sendNotebookEvent(sessionGeneration, {
+      event: "marimo",
+      data: { marimo: notification },
+    });
+    return [notification, sessionGeneration];
   }
 
   async function initializeSingleApp() {
@@ -300,7 +301,7 @@ describe("IslandsPyodideBridge", () => {
   describe("app lifecycle", () => {
     it("ignores messages after a session yields ownership", async () => {
       const consumeMessage = vi.fn();
-      bridge.consumeMessages(consumeMessage);
+      bridge.consumeNotifications(consumeMessage);
       mockSingleApp();
       await bridge.initializeApps();
 
@@ -319,10 +320,7 @@ describe("IslandsPyodideBridge", () => {
       sendKernelMessage(1);
       const currentMessage = sendKernelMessage(2);
 
-      expect(consumeMessage.mock.calls).toEqual([
-        [firstMessage],
-        [currentMessage],
-      ]);
+      expect(consumeMessage.mock.calls).toEqual([firstMessage, currentMessage]);
 
       finishReplacement();
       await replacement;
@@ -330,7 +328,7 @@ describe("IslandsPyodideBridge", () => {
 
     it("restores message ownership when replacement fails", async () => {
       const consumeMessage = vi.fn();
-      bridge.consumeMessages(consumeMessage);
+      bridge.consumeNotifications(consumeMessage);
       mockSingleApp();
       await bridge.initializeApps();
       mockReplaceSessionRequest.mockRejectedValueOnce(
@@ -344,12 +342,12 @@ describe("IslandsPyodideBridge", () => {
       const previousMessage = sendKernelMessage(1);
       sendKernelMessage(2);
 
-      expect(consumeMessage.mock.calls).toEqual([[previousMessage]]);
+      expect(consumeMessage.mock.calls).toEqual([previousMessage]);
     });
 
     it("ignores messages while their session is stopping", async () => {
       const consumeMessage = vi.fn();
-      bridge.consumeMessages(consumeMessage);
+      bridge.consumeNotifications(consumeMessage);
       mockSingleApp();
       await bridge.initializeApps();
       let finishStop!: () => void;
@@ -373,7 +371,7 @@ describe("IslandsPyodideBridge", () => {
 
     it("forwards messages from every app in a multi-app document", async () => {
       const consumeMessage = vi.fn();
-      bridge.consumeMessages(consumeMessage);
+      bridge.consumeNotifications(consumeMessage);
       mockParseMarimoIslandApps.mockReturnValue([app(), app("app-2")]);
       mockCreateMarimoFile
         .mockReturnValueOnce("generated app 1")
@@ -385,9 +383,58 @@ describe("IslandsPyodideBridge", () => {
       const secondMessage = sendKernelMessage(2);
       sendKernelMessage(3);
 
-      expect(consumeMessage.mock.calls).toEqual([
-        [firstMessage],
-        [secondMessage],
+      expect(consumeMessage.mock.calls).toEqual([firstMessage, secondMessage]);
+    });
+
+    it("delivers only the notebook stream's marimo notifications", async () => {
+      const consumeMessage = vi.fn();
+      bridge.consumeNotifications(consumeMessage);
+      await initializeSingleApp();
+      sendNotebookEvent(1, { event: "ready", id: "3", data: {} });
+      const notification = sendKernelMessage(1);
+
+      expect(consumeMessage.mock.calls).toEqual([notification]);
+    });
+
+    it("forwards the managed session's host events to the notebook stream", async () => {
+      await initializeSingleApp();
+      const handlers = { open: vi.fn(), message: vi.fn() };
+      void bridge.notebookBackend.transport(
+        { signal: new AbortController().signal, cursor: null },
+        handlers,
+      );
+      expect(handlers.open).toHaveBeenCalledOnce();
+
+      const cell = { event: "cell", id: "1", data: { cell: { id: "c1" } } };
+      sendNotebookEvent(1, cell);
+      sendNotebookEvent(2, { event: "ready", id: "2", data: {} });
+      sendKernelMessage(1);
+
+      expect(handlers.message.mock.calls).toEqual([[cell]]);
+      expect(() => bridge.notebookBackend.notebook("n1")).toThrow(
+        "Islands do not edit",
+      );
+    });
+
+    it("resets the notebook stream when the managed session is replaced or stopped", async () => {
+      await initializeSingleApp();
+      const handlers = { open: vi.fn(), message: vi.fn() };
+      void bridge.notebookBackend.transport(
+        { signal: new AbortController().signal, cursor: null },
+        handlers,
+      );
+
+      mockSingleApp("generated app 2");
+      await bridge.initializeApps();
+      const ready = { event: "ready", id: "2", data: {} };
+      sendNotebookEvent(1, { event: "ready", id: "1", data: {} });
+      sendNotebookEvent(2, ready);
+      await bridge.stopSession("app-1");
+
+      expect(handlers.message.mock.calls).toEqual([
+        [{ event: "reset", data: {} }],
+        [ready],
+        [{ event: "reset", data: {} }],
       ]);
     });
 

package/src/core/islands/bootstrap.ts

@@ -9,11 +9,7 @@ import {
 } from "@/plugins/impl/anywidget/registry";
 import { initializePlugins } from "@/plugins/plugins";
 import { Functions } from "@/utils/functions";
-import {
-  safeExtractSetUIElementMessageBuffers,
-  type JsonString,
-} from "@/utils/json/base64";
-import { jsonParseWithSpecialChar } from "@/utils/json/json-parser";
+import { safeExtractSetUIElementMessageBuffers } from "@/utils/json/base64";
 import { Logger } from "@/utils/Logger";
 import { logNever } from "@/utils/assertNever";
 import {
@@ -21,19 +17,20 @@ import {
   notebookAtom,
   notebookReducer,
 } from "../cells/cells";
+import { followDocument } from "../cells/document";
+import { notebookCells } from "../cells/utils";
+import { createDocumentModel } from "../notebook/document";
 import { defineCustomElement } from "../dom/defineCustomElement";
 import { MarimoValueInputEvent } from "../dom/events";
 import { UI_ELEMENT_REGISTRY } from "../dom/uiregistry";
 import { FUNCTIONS_REGISTRY } from "../functions/FunctionRegistry";
+import { createNotebookReplica } from "../notebook/replica";
 import {
   handleCellNotificationeration,
   handleKernelReady,
   handleRemoveUIElements,
 } from "../kernel/handlers";
-import type {
-  NotificationMessage,
-  NotificationMessageType,
-} from "../kernel/messages";
+import type { NotificationMessage } from "../kernel/messages";
 import { queryParamHandlers } from "../kernel/queryParamHandlers";
 import { RuntimeState } from "../kernel/RuntimeState";
 import { initialModeAtom } from "../mode";
@@ -126,19 +123,26 @@ export async function initializeIslands(
     }
   });
 
-  // Wire up kernel message handling
-  // The wire format is {"op": "...", "data": {"op": "...", ...}} — both
-  // the envelope and the payload carry the op. The bridge types the message
-  // as NotificationPayload (just {data}), but the actual wire format
-  // includes the outer op too.
-  // Worker metadata identifies the session that emitted the message.
-  bridge.consumeMessages(({ message, sessionGeneration }) => {
-    handleMessage(
-      message as unknown as JsonString<IslandsNotificationMessage>,
-      actions,
-      store,
-      sessionGeneration,
-    );
+  // Each notification names the session that sent it.
+  bridge.consumeNotifications((notification, sessionGeneration) => {
+    handleNotification(notification, actions, store, sessionGeneration);
+  });
+
+  // The notebook's cells follow the managed session's host document.
+  const replica = createNotebookReplica(store);
+  replica.connect(bridge.notebookBackend, { onNotification: Functions.NOOP });
+  const documentModel = createDocumentModel({
+    store,
+    documentAtom: replica.documentAtom,
+    backendAtom: replica.backendAtom,
+  });
+  const lifetime = new AbortController();
+  documentModel.start({ signal: lifetime.signal });
+  followDocument({
+    store,
+    model: documentModel,
+    reconcile: actions.reconcileDocument,
+    signal: lifetime.signal,
   });
 
   // Track first user interaction before initialization completes
@@ -154,27 +158,15 @@ export async function initializeIslands(
   RuntimeState.INSTANCE.start(bridge.sendComponentValues);
 }
 
-type IslandsNotificationMessage = {
-  [K in NotificationMessageType]: {
-    data: Extract<NotificationMessage, { op: K }>;
-    op: K;
-  };
-}[NotificationMessageType];
-
-/**
- * Handles a single message from the kernel.
- *
- * Wire format from Python: {"op": "<name>", "data": {"op": "<name>", ...}}
- */
+/** Handles a single notification from the kernel. */
 // oxlint-disable-next-line marimo/prefer-object-params
-function handleMessage(
-  message: JsonString<IslandsNotificationMessage>,
+function handleNotification(
+  msg: NotificationMessage,
   actions: NotebookActions,
   store: Store,
   sessionGeneration: number,
 ): void {
   try {
-    const msg = jsonParseWithSpecialChar(message);
     Logger.debug("Islands received message:", msg.op);
 
     switch (msg.op) {
@@ -223,11 +215,12 @@ function handleMessage(
         return;
 
       case "kernel-ready":
-        handleKernelReady(msg.data, {
+        handleKernelReady(msg, {
           autoInstantiate: true,
-          setCells: actions.setCells,
+          cells: notebookCells(store.get(notebookAtom)),
+          setLastRun: actions.setLastRun,
+          resetCellRuntimes: actions.resetCellRuntimes,
           setLayoutData: Functions.NOOP,
-          setAppConfig: Functions.NOOP,
           setCapabilities: Functions.NOOP,
           setKernelState: Functions.NOOP,
           onError: Logger.error,
@@ -239,47 +232,47 @@ function handleMessage(
 
       case "send-ui-element-message":
         UI_ELEMENT_REGISTRY.broadcastMessage(
-          msg.data.ui_element,
-          msg.data.message,
-          safeExtractSetUIElementMessageBuffers(msg.data),
+          msg.ui_element,
+          msg.message,
+          safeExtractSetUIElementMessageBuffers(msg),
         );
         return;
 
       case "remove-ui-elements":
-        handleRemoveUIElements(msg.data);
+        handleRemoveUIElements(msg);
         return;
 
       case "function-call-result":
-        FUNCTIONS_REGISTRY.resolve(msg.data.function_call_id, msg.data);
+        FUNCTIONS_REGISTRY.resolve(msg.function_call_id, msg);
         return;
 
       case "cell-op":
-        handleCellNotificationeration(msg.data, actions.handleCellMessage);
+        handleCellNotificationeration(msg, actions.handleCellMessage);
         return;
 
       case "alert":
         toast({
-          title: msg.data.title,
-          description: renderHTML({ html: msg.data.description }),
-          variant: msg.data.variant,
+          title: msg.title,
+          description: renderHTML({ html: msg.description }),
+          variant: msg.variant,
         });
         return;
 
       case "query-params-append":
-        queryParamHandlers.append(msg.data);
+        queryParamHandlers.append(msg);
         return;
       case "query-params-set":
-        queryParamHandlers.set(msg.data);
+        queryParamHandlers.set(msg);
         return;
       case "query-params-delete":
-        queryParamHandlers.delete(msg.data);
+        queryParamHandlers.delete(msg);
         return;
       case "query-params-clear":
         queryParamHandlers.clear();
         return;
 
       case "model-lifecycle":
-        handleWidgetMessage(WIDGET_REGISTRY, msg.data);
+        handleWidgetMessage(WIDGET_REGISTRY, msg);
         return;
 
       case "consumer-capabilities":