@marimo-team/islands 0.25.3-dev28 → 0.25.3-dev281

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,26 @@ 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 { 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 +125,27 @@ 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 });
+  followDocument({
+    store,
+    model: documentModel,
+    reconcile: actions.reconcileDocument,
+    signal: lifetime.signal,
   });
 
   // Track first user interaction before initialization completes
@@ -154,27 +161,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,13 +218,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 +233,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":

package/src/core/islands/bridge.ts

@@ -1,19 +1,23 @@
 /* Copyright 2026 Marimo. All rights reserved. */
 /* oxlint-disable typescript/no-explicit-any */
 
+import type { NotebookTransport } from "@marimo-team/client";
 import { getWorkerRPC } from "@/core/wasm/rpc";
+import { Arrays } from "@/utils/arrays";
 import { Deferred } from "@/utils/Deferred";
 import { throwNotImplemented } from "@/utils/functions";
 import { Logger } from "@/utils/Logger";
+import { jsonParseWithSpecialChar } from "@/utils/json/json-parser";
 import { generateUUID } from "@/utils/uuid";
 import { initialNotebookState, notebookAtom } from "../cells/cells";
-import type { ClientCommand } from "../kernel/messages";
+import type { NotebookBackend } from "../notebook/replica";
+import type { ClientCommand, NotificationMessage } from "../kernel/messages";
 import type { EditRequests, RunRequests } from "../network/types";
 import { store as defaultStore } from "../state/jotai";
 import { getMarimoExportContext } from "../static/export-context";
 import { createMarimoFile, parseMarimoIslandApps } from "./parse";
 import { islandsInitializedAtom, islandsPendingInitialRunsAtom } from "./state";
-import type { IslandsKernelMessage, WorkerSchema } from "./worker/worker";
+import type { WorkerSchema } from "./worker/worker";
 import type { WorkerFactory } from "./worker-factory";
 import { DefaultWorkerFactory } from "./worker-factory";
 
@@ -52,14 +56,14 @@ interface AppSession {
  * @example
  * ```ts
  * const bridge = new IslandsPyodideBridge();
- * bridge.consumeMessages(message => console.log(message));
+ * bridge.consumeNotifications((notification) => console.log(notification.op));
  * await bridge.initializeApps();
  * ```
  */
 export class IslandsPyodideBridge implements RunRequests, EditRequests {
   private rpc: ReturnType<typeof getWorkerRPC<WorkerSchema>>;
-  private messageConsumer:
-    | ((message: IslandsKernelMessage) => void)
+  private notificationConsumer:
+    | ((notification: NotificationMessage, sessionGeneration: number) => void)
     | undefined;
   private readonly store: typeof defaultStore;
   private readonly root: Document | Element;
@@ -69,9 +73,37 @@ export class IslandsPyodideBridge implements RunRequests, EditRequests {
   private readonly activeSessionGenerations = new Set<number>();
   private appTransition = Promise.resolve();
   private workerReady = new Deferred<void>();
+  /** Where the managed session's host events go while a replica follows them. */
+  private stream: Parameters<NotebookTransport>[1] | undefined;
 
   public initialized = new Deferred<void>();
 
+  /**
+   * The managed session's host stream, for a replica to follow. Replacing
+   * or stopping the session resets the stream, and the next session's
+   * snapshot follows.
+   */
+  public readonly notebookBackend: NotebookBackend = {
+    transport: ({ signal }, handlers) =>
+      new Promise((resolve) => {
+        this.stream = handlers;
+        handlers.open();
+        signal.addEventListener(
+          "abort",
+          () => {
+            if (this.stream === handlers) {
+              this.stream = undefined;
+            }
+            resolve();
+          },
+          { once: true },
+        );
+      }),
+    notebook: () => {
+      throw new Error("Islands do not edit notebooks");
+    },
+  };
+
   public constructor(config: IslandsBridgeConfig = {}) {
     this.store = config.store || defaultStore;
     this.root = config.root || document;
@@ -110,14 +142,35 @@ export class IslandsPyodideBridge implements RunRequests, EditRequests {
       },
     );
 
-    this.rpc.addMessageListener("kernelMessage", (message) => {
-      // A stopped session can still have messages in transit after the next
-      // app takes ownership of the shared frontend state.
-      if (!this.activeSessionGenerations.has(message.sessionGeneration)) {
-        return;
-      }
-      this.messageConsumer?.(message);
-    });
+    this.rpc.addMessageListener(
+      "notebookEvent",
+      ({ event, sessionGeneration }) => {
+        // A stopped session can still have events in transit after the next
+        // app takes ownership of the shared frontend state.
+        if (!this.activeSessionGenerations.has(sessionGeneration)) {
+          return;
+        }
+        const frame = jsonParseWithSpecialChar(event) as {
+          event: string;
+          data: { marimo?: NotificationMessage };
+        };
+        if (frame.event === "marimo") {
+          if (frame.data.marimo) {
+            this.notificationConsumer?.(frame.data.marimo, sessionGeneration);
+          }
+          return;
+        }
+        // The replica follows the managed session's host document.
+        if (sessionGeneration === this.session?.sessionGeneration) {
+          this.stream?.message(frame);
+        }
+      },
+    );
+  }
+
+  /** The managed session is gone, so its stream starts over. */
+  private resetNotebookStream(): void {
+    this.stream?.message({ event: "reset", data: {} });
   }
 
   public async initializeApps(): Promise<void> {
@@ -174,6 +227,7 @@ export class IslandsPyodideBridge implements RunRequests, EditRequests {
       const replacesSession = managesSingleApp && previousSession !== undefined;
       this.activeSessionGenerations.add(request.sessionGeneration);
       if (replacesSession) {
+        this.resetNotebookStream();
         this.store.set(notebookAtom, initialNotebookState());
       }
       if (managesSingleApp || !previousSession) {
@@ -225,6 +279,7 @@ export class IslandsPyodideBridge implements RunRequests, EditRequests {
       this.store.set(islandsPendingInitialRunsAtom, null);
       this.session = undefined;
       this.sessionReady = new Deferred<AppSession>();
+      this.resetNotebookStream();
       this.store.set(notebookAtom, initialNotebookState());
     });
   }
@@ -246,14 +301,14 @@ export class IslandsPyodideBridge implements RunRequests, EditRequests {
     await this.rpc.proxy.request.startSession(opts);
   }
 
-  /**
-   * Sets up a consumer for kernel messages
-   */
-  public consumeMessages(
-    consumer: (message: IslandsKernelMessage) => void,
+  /** Receive each running session's marimo notifications. */
+  public consumeNotifications(
+    consumer: (
+      notification: NotificationMessage,
+      sessionGeneration: number,
+    ) => void,
   ): void {
-    this.messageConsumer = consumer;
-    this.rpc.proxy.send.consumerReady({});
+    this.notificationConsumer = consumer;
   }
 
   // ============================================================================
@@ -308,7 +363,14 @@ export class IslandsPyodideBridge implements RunRequests, EditRequests {
     await this.putControlRequest(
       {
         type: "execute-cells",
-        ...request,
+        request: request.request,
+        executionRequests: Arrays.zip(request.cellIds, request.codes).map(
+          ([cellId, code]) => ({
+            type: "execute-cell",
+            cellId,
+            code,
+          }),
+        ),
       },
       session,
     );
@@ -322,21 +384,18 @@ export class IslandsPyodideBridge implements RunRequests, EditRequests {
   public getUsageStats = throwNotImplemented;
   public getEnvironmentInfo = throwNotImplemented;
   public sendRename = throwNotImplemented;
-  public sendSave = throwNotImplemented;
+  public saveLayout = throwNotImplemented;
   public sendCopy = throwNotImplemented;
   public sendRunScratchpad = throwNotImplemented;
   public sendStdin = throwNotImplemented;
   public sendInterrupt = throwNotImplemented;
   public sendShutdown = throwNotImplemented;
   public sendFormat = throwNotImplemented;
-  public sendDeleteCell = throwNotImplemented;
   public sendInstallMissingPackages = throwNotImplemented;
   public sendCodeCompletionRequest = throwNotImplemented;
   public saveUserConfig = throwNotImplemented;
   public saveAppConfig = throwNotImplemented;
-  public saveCellConfig = throwNotImplemented;
   public sendRestart = throwNotImplemented;
-  public sendDocumentTransaction = throwNotImplemented;
   public readCode = throwNotImplemented;
   public readSnippets = throwNotImplemented;
   public previewDatasetColumn = throwNotImplemented;

package/src/core/islands/components/__tests__/IslandControls.test.tsx

@@ -2,6 +2,7 @@
 
 import { fireEvent, render, screen } from "@testing-library/react";
 import { beforeEach, describe, expect, it, vi } from "vitest";
+import { MockRequestClient } from "@/__mocks__/requests";
 import { TooltipProvider } from "@/components/ui/tooltip";
 import type { CellId } from "@/core/cells/ids";
 import * as requestsModule from "@/core/network/requests";
@@ -18,7 +19,8 @@ vi.mock("@/utils/copy", () => ({
 }));
 
 describe("IslandControls", () => {
-  const mockSendRun = vi.fn();
+  const mockSendRun =
+    vi.fn<ReturnType<typeof requestsModule.useRequestClient>["sendRun"]>();
   const mockCopyToClipboard = vi.fn();
   const mockCodeCallback = vi.fn(() => "print('test code')");
   const cellId = "test-cell-id" as CellId;
@@ -30,11 +32,11 @@ describe("IslandControls", () => {
 
   beforeEach(() => {
     vi.clearAllMocks();
-    mockSendRun.mockResolvedValue(undefined);
+    mockSendRun.mockResolvedValue(null);
 
-    vi.spyOn(requestsModule, "useRequestClient").mockReturnValue({
-      sendRun: mockSendRun,
-    } as any);
+    vi.spyOn(requestsModule, "useRequestClient").mockReturnValue(
+      MockRequestClient.create({ sendRun: mockSendRun }),
+    );
 
     vi.spyOn(copyModule, "copyToClipboard").mockImplementation(
       mockCopyToClipboard,

package/src/core/islands/components/__tests__/web-components.test.tsx

@@ -9,6 +9,7 @@ import {
   describe,
   expect,
   it,
+  vi,
 } from "vitest";
 import { cellId } from "@/__tests__/branded";
 import { MockNotebook } from "@/__mocks__/notebook";
@@ -22,6 +23,7 @@ import {
 import { islandsPendingInitialRunsAtom } from "@/core/islands/state";
 import { requestClientAtom } from "@/core/network/requests";
 import { store } from "@/core/state/jotai";
+import { buildIslandHTML } from "../../__tests__/test-utils";
 import { parseMarimoIslandApps } from "../../parse";
 import { MarimoIslandElement } from "../web-components";
 
@@ -45,9 +47,55 @@ beforeEach(() => {
 afterEach(async () => {
   await actAndFlush(() => document.body.replaceChildren());
   store.set(requestClientAtom, null);
+  document.body.removeAttribute("style");
+  vi.unstubAllEnvs();
 });
 
 describe("MarimoIslandElement lifecycle", () => {
+  it.each([true, false])(
+    "synchronizes the theme of a reactive=%s island across reconnects",
+    async (reactive) => {
+      vi.stubEnv("VITE_MARIMO_ISLANDS", "true");
+      if (reactive) {
+        setNotebook("runtime-cell", "runtime output");
+      }
+      const { container, island } = createThemeIsland(reactive);
+      document.body.style.colorScheme = "dark";
+      await actAndFlush(() => document.body.append(container));
+      expect(island).toHaveClass("marimo", "dark");
+      await actAndFlush(() => {
+        document.body.style.colorScheme = "light";
+      });
+      expect(island).toHaveClass("marimo", "light");
+      expect(island).not.toHaveClass("dark");
+      await actAndFlush(() => island.remove());
+      await actAndFlush(() => {
+        document.body.style.colorScheme = "dark";
+      });
+      expect(island).toHaveClass("light");
+      await actAndFlush(() => container.append(island));
+      expect(island).toHaveClass("dark");
+    },
+  );
+
+  it("can disconnect before its scheduled render and reconnect with the current theme", async () => {
+    vi.stubEnv("VITE_MARIMO_ISLANDS", "true");
+    const { container, island } = createThemeIsland(false);
+    document.body.style.colorScheme = "dark";
+    await actAndFlush(() => {
+      document.body.append(container);
+      island.remove();
+    });
+    expect(island.querySelector("marimo-cell-output")).not.toBeNull();
+    await actAndFlush(() => {
+      document.body.style.colorScheme = "light";
+    });
+    expect(island).toHaveClass("dark");
+    await actAndFlush(() => container.append(island));
+    expect(island).toHaveClass("light");
+    expect(island.textContent).toContain("static output");
+  });
+
   it("keeps static output until the initial run completes", async () => {
     store.set(islandsPendingInitialRunsAtom, new Set([1]));
     setNotebook("runtime-cell", "runtime output");
@@ -182,6 +230,21 @@ describe("MarimoIslandElement lifecycle", () => {
   });
 });
 
+function createThemeIsland(reactive: boolean) {
+  const container = document.createElement("div");
+  container.innerHTML = buildIslandHTML([
+    {
+      appId: "app-1",
+      cellId: "runtime-cell",
+      reactive,
+      code: "x = 1",
+      output: "<pre>static output</pre>",
+    },
+  ]);
+  const island = container.querySelector<HTMLElement>(ISLAND_TAG_NAMES.ISLAND)!;
+  return { container, island };
+}
+
 function setNotebook(id: string, output?: string): void {
   const runtimeCellId = cellId(id);
   store.set(

package/src/core/islands/components/output-wrapper.tsx

@@ -13,7 +13,7 @@ import { OutputRenderer } from "@/components/editor/Output";
 import { type NotebookState, notebookAtom } from "@/core/cells/cells";
 import type { CellId } from "@/core/cells/ids";
 import { isOutputEmpty } from "@/core/cells/outputs";
-import type { CellRuntimeState } from "@/core/cells/types";
+import { type CellRunStatus, useCellExecution } from "@/core/executions";
 import { ISLAND_TAG_NAMES } from "../constants";
 import { islandsHydratedAtom } from "../state";
 import { IslandControls } from "./IslandControls";
@@ -69,7 +69,7 @@ export const MarimoOutputWrapper: React.FC<MarimoOutputWrapperProps> = ({
 
   // Sync cell status to the host <marimo-island> element as a data attribute
   // so downstream consumers can style based on status (e.g. [data-status="running"])
-  const status = runtime?.status ?? "idle";
+  const { status } = useCellExecution(cellId);
   useSyncStatusToIsland(wrapperRef, status);
 
   // If no runtime yet, show static content
@@ -90,7 +90,7 @@ export const MarimoOutputWrapper: React.FC<MarimoOutputWrapperProps> = ({
   return (
     <div ref={wrapperRef} className="relative min-h-6">
       <OutputRenderer message={runtime.output} />
-      <RunningIndicator state={runtime} />
+      <RunningIndicator status={status} />
       <IslandControls
         cellId={cellId}
         codeCallback={codeCallback}
@@ -118,8 +118,8 @@ function useSyncStatusToIsland(
 /**
  * Shows a spinning indicator when the cell is running
  */
-const RunningIndicator: React.FC<{ state: CellRuntimeState }> = ({ state }) => {
-  if (state.status === "running") {
+const RunningIndicator: React.FC<{ status: CellRunStatus }> = ({ status }) => {
+  if (status === "running") {
     return (
       <DelayRender>
         <div className="absolute top-1 right-1">

package/src/core/islands/components/web-components.tsx

@@ -10,6 +10,7 @@ import { OBJECT_ID_ATTR } from "@/core/dom/ui-element-constants";
 import { UI_ELEMENT_REGISTRY } from "@/core/dom/uiregistry";
 import { LocaleProvider } from "@/core/i18n/locale-provider";
 import { renderHTML } from "@/plugins/core/RenderHTML";
+import { resolvedThemeAtom } from "@/theme/useTheme";
 import { invariant } from "@/utils/invariant";
 import type { CellId, UIElementId } from "../../cells/ids";
 import { store } from "../../state/jotai";
@@ -44,6 +45,7 @@ export class MarimoIslandElement extends HTMLElement {
   private runtimeAppId?: string;
   private runtimeCellId?: CellId;
   private unsubscribeCellIds?: () => void;
+  private unsubscribeTheme?: () => void;
 
   public static readonly tagName = ISLAND_TAG_NAMES.ISLAND;
   public static readonly outputTagName = ISLAND_TAG_NAMES.CELL_OUTPUT;
@@ -104,7 +106,7 @@ export class MarimoIslandElement extends HTMLElement {
    */
   private getCellIdFromIndex(idx: number): CellId | undefined {
     const { cellIds } = store.get(notebookAtom);
-    return cellIds.inOrderIds.at(idx);
+    return cellIds.at(idx);
   }
 
   /**
@@ -132,6 +134,13 @@ export class MarimoIslandElement extends HTMLElement {
       ? this.cellId
       : undefined;
     this.syncCellIdsSubscription();
+    const syncTheme = () => {
+      const theme = store.get(resolvedThemeAtom);
+      this.classList.toggle("dark", theme === "dark");
+      this.classList.toggle("light", theme === "light");
+    };
+    this.unsubscribeTheme = store.sub(resolvedThemeAtom, syncTheme);
+    syncTheme();
     queueMicrotask(() => {
       // Guard against disconnect between connectedCallback and microtask
       if (
@@ -297,6 +306,8 @@ export class MarimoIslandElement extends HTMLElement {
    * Cleanup when element is removed from DOM
    */
   public disconnectedCallback(): void {
+    this.unsubscribeTheme?.();
+    this.unsubscribeTheme = undefined;
     this.removeEventListener(
       ISLAND_SOURCE_CHANGED_EVENT,
       this.handleSourceChanged,