@marimo-team/islands 0.25.3-dev34 → 0.25.3-dev349

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/hotkeys/hotkeys.ts

@@ -221,7 +221,8 @@ const DEFAULT_HOT_KEY = {
     key: "Mod-h",
   },
   "cell.aiCompletion": {
-    name: "AI completion",
+    name: "Refactor cell",
+    additionalKeywords: ["AI", "AI completion"],
     group: "Editing",
     key: "Mod-Shift-e",
   },

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

@@ -2,15 +2,9 @@
 
 import type { components } from "@marimo-team/marimo-api";
 import { beforeEach, describe, expect, it, vi } from "vitest";
-import {
-  cellId,
-  requestId,
-  uiElementId,
-  widgetModelId,
-} from "@/__tests__/branded";
+import { requestId, uiElementId, widgetModelId } 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 {
@@ -51,6 +45,7 @@ vi.stubGlobal("URL", MockURL);
 const {
   mockBridge,
   mockLoadPackages,
+  mockNotebookCall,
   mockReplaceSessionRequest,
   mockStopSessionRequest,
   mockStartSessionRequest,
@@ -62,6 +57,7 @@ const {
 } = vi.hoisted(() => ({
   mockBridge: vi.fn(),
   mockLoadPackages: vi.fn(),
+  mockNotebookCall: vi.fn(),
   mockReplaceSessionRequest: vi.fn(),
   mockStopSessionRequest: vi.fn(),
   mockStartSessionRequest: vi.fn(),
@@ -85,13 +81,11 @@ vi.mock("@/core/wasm/rpc", () => ({
       request: {
         bridge: mockBridge,
         loadPackages: mockLoadPackages,
+        notebookCall: mockNotebookCall,
         replaceSession: mockReplaceSessionRequest,
         startSession: mockStartSessionRequest,
         stopSession: mockStopSessionRequest,
       },
-      send: {
-        consumerReady: vi.fn(),
-      },
     },
     addMessageListener: vi.fn(
       (name: string, listener: (payload: never) => void) => {
@@ -177,17 +171,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 +299,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 +318,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 +326,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 +340,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 +369,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 +381,55 @@ 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]]);
+    });
+
+    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: {} }],
       ]);
     });
 
@@ -707,49 +749,47 @@ describe("IslandsPyodideBridge", () => {
     });
   });
 
-  describe("sendRun", () => {
+  describe("notebookBackend.notebook", () => {
     beforeEach(initializeSingleApp);
 
-    it("should include type field in control request", async () => {
-      const request = {
-        cellIds: [cellId("cell-1"), cellId("cell-2")],
-        codes: ["print('hello')", "print('world')"],
-      };
-
-      await bridge.sendRun(request);
-
-      expect(mockBridge).toHaveBeenCalledWith({
-        appId: "app-1",
-        functionName: "put_control_request",
-        payload: {
-          type: "execute-cells",
-          cellIds: ["cell-1", "cell-2"],
-          codes: ["print('hello')", "print('world')"],
+    it("sends notebook commands to the managed session's host", async () => {
+      mockNotebookCall.mockResolvedValue({
+        status: 202,
+        body: {
+          id: "run-1",
+          notebook_id: "n1",
+          runtime_id: "r1",
+          cells: ["c1"],
+          status: "queued",
+          outcome: null,
+          started_at: null,
+          completed_at: null,
         },
-        sessionGeneration: 1,
       });
-    });
-
-    it("should call loadPackages before putControlRequest", async () => {
-      const request = {
-        cellIds: [cellId("cell-1")],
-        codes: ["import pandas"],
-      };
 
-      await bridge.sendRun(request);
+      const result = await bridge.notebookBackend
+        .notebook("n1")
+        .run({ cells: [{ id: "c1" }], idempotencyKey: "key-1" });
 
-      // Verify loadPackages was called with joined codes
-      expect(mockLoadPackages).toHaveBeenCalledWith({
+      expect(result.ok).toBe(true);
+      expect(mockNotebookCall).toHaveBeenCalledWith({
         appId: "app-1",
-        code: "import pandas",
         sessionGeneration: 1,
+        call: expect.objectContaining({
+          operation: "runtime.run",
+          params: { notebook_id: "n1" },
+          headers: { "idempotency-key": "key-1" },
+          body: { cells: [{ id: "c1" }] },
+        }),
       });
+    });
 
-      // Verify order: loadPackages should be called before bridge
-      const loadPackagesCallOrder =
-        mockLoadPackages.mock.invocationCallOrder[0];
-      const bridgeCallOrder = mockBridge.mock.invocationCallOrder[0];
-      expect(loadPackagesCallOrder).toBeLessThan(bridgeCallOrder);
+    it("has no commands without an active session", async () => {
+      await bridge.stopSession();
+
+      expect(() => bridge.notebookBackend.notebook("n1")).toThrow(
+        "No island session is active",
+      );
     });
   });
 
@@ -799,7 +839,6 @@ describe("IslandsPyodideBridge", () => {
         functionName: "",
         args: {},
       });
-      await bridge.sendRun({ cellIds: [], codes: [] });
       await bridge.sendModelValue({
         modelId: widgetModelId(""),
         message: { method: "update", state: {}, bufferPaths: [] },

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,23 +17,29 @@ 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 { observeCellMessage } from "../executions";
 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";
 import { requestClientAtom } from "../network/requests";
+import {
+  activeDocumentModelAtom,
+  activeNotebookReplicaAtom,
+} from "../session/active-notebook";
 import { store as defaultStore } from "../state/jotai";
 import type { IslandsPyodideBridge } from "./bridge";
 import { MarimoIslandElement } from "./components/web-components";
@@ -126,19 +128,28 @@ 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 });
+  store.set(activeNotebookReplicaAtom, replica);
+  const documentModel = createDocumentModel({
+    store,
+    documentAtom: replica.documentAtom,
+    backendAtom: replica.backendAtom,
+  });
+  const lifetime = new AbortController();
+  documentModel.start({ signal: lifetime.signal });
+  store.set(activeDocumentModelAtom, documentModel);
+  followDocument({
+    store,
+    model: documentModel,
+    reconcile: actions.reconcileDocument,
+    signal: lifetime.signal,
   });
 
   // Track first user interaction before initialization completes
@@ -154,27 +165,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) {
@@ -204,7 +203,6 @@ function handleMessage(
       case "startup-logs":
       case "startup-progress":
       case "interrupted":
-      case "reconnected":
       case "cache-cleared":
       case "cache-info":
       case "kernel-startup-error":
@@ -223,13 +221,12 @@ function handleMessage(
         return;
 
       case "kernel-ready":
-        handleKernelReady(msg.data, {
+        handleKernelReady(msg, {
           autoInstantiate: true,
-          setCells: actions.setCells,
+          cells: notebookCells(store.get(notebookAtom)),
+          resetCellRuntimes: actions.resetCellRuntimes,
           setLayoutData: Functions.NOOP,
-          setAppConfig: Functions.NOOP,
           setCapabilities: Functions.NOOP,
-          setKernelState: Functions.NOOP,
           onError: Logger.error,
         });
         if (!window.customElements?.get(ISLAND_TAG_NAMES.ISLAND)) {
@@ -239,47 +236,48 @@ 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);
+        observeCellMessage(store, msg);
+        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":