@marimo-team/islands 0.25.3-dev19 → 0.25.3-dev191

package/src/core/islands/worker/controller.ts

@@ -1,11 +1,9 @@
 /* Copyright 2026 Marimo. All rights reserved. */
 import type { PyodideInterface } from "pyodide";
 import { defaultUserConfig } from "@/core/config/config-schema";
-import type { NotificationPayload } from "@/core/kernel/messages";
 import { DefaultWasmController } from "@/core/wasm/worker/bootstrap";
 import { WasmFileSystem } from "@/core/wasm/worker/fs";
-import type { SerializedBridge } from "@/core/wasm/worker/types";
-import type { JsonString } from "@/utils/json/base64";
+import type { WasmSession } from "@/core/wasm/worker/types";
 import { Logger } from "../../../utils/Logger";
 
 export class ReadonlyWasmController extends DefaultWasmController {
@@ -37,19 +35,15 @@ export class ReadonlyWasmController extends DefaultWasmController {
     return { code, filename };
   }
 
-  public override async startSession(opts: {
+  public override async prepareSession(opts: {
     code: string;
     filename: string | null;
-    onMessage: (message: JsonString<NotificationPayload>) => void;
-  }): Promise<SerializedBridge> {
-    const bridge = super.startSession({
+  }): Promise<WasmSession> {
+    return super.prepareSession({
       queryParameters: {},
       code: opts.code,
       filename: opts.filename,
-      onMessage: opts.onMessage,
       userConfig: defaultUserConfig(),
     });
-
-    return bridge;
   }
 }

package/src/core/islands/worker/worker.tsx

@@ -1,21 +1,17 @@
 /* Copyright 2026 Marimo. All rights reserved. */
 
 import type { PyodideInterface } from "pyodide";
-import type { PyProxy } from "pyodide/ffi";
 import {
   createRPC,
   createRPCRequestHandler,
   createWorkerParentTransport,
   type RPCSchema,
 } from "rpc-anywhere";
-import type { NotificationPayload } from "@/core/kernel/messages";
 import type { ParentSchema } from "@/core/wasm/rpc";
 import { shouldLoadDuckDBPackages } from "@/core/wasm/utils";
 import { TRANSPORT_ID } from "@/core/wasm/worker/constants";
 import { getPyodideVersion } from "@/core/wasm/worker/getPyodideVersion";
-import { MessageBuffer } from "@/core/wasm/worker/message-buffer";
 import type { RawBridge, SerializedBridge } from "@/core/wasm/worker/types";
-import type { JsonString } from "@/utils/json/base64";
 import { Logger } from "@/utils/Logger";
 import { prettyError } from "../../../utils/errors";
 import { invariant } from "../../../utils/invariant";
@@ -45,14 +41,11 @@ async function loadPyodideAndPackages() {
 }
 
 const pyodideReadyPromise = loadPyodideAndPackages();
-export interface IslandsKernelMessage {
-  message: JsonString<NotificationPayload>;
+/** An event of one session's notebook stream. */
+export interface IslandsNotebookEvent {
+  event: string;
   sessionGeneration: number;
 }
-
-const messageBuffer = new MessageBuffer<IslandsKernelMessage>((payload) => {
-  rpc.send.kernelMessage(payload);
-});
 interface SessionRequest {
   code: string;
   appId: string;
@@ -79,21 +72,22 @@ async function startSession(
       code: opts.code,
       filename: `app-${opts.appId}.py`,
     });
-    const nextBridge = await self.controller.startSession({
-      ...notebook,
-      onMessage: (message) => {
-        messageBuffer.push({
-          message,
-          sessionGeneration: opts.sessionGeneration,
-        });
-      },
+    const session = await self.controller.prepareSession(notebook);
+    // Follow before the session instantiates. Each subscription begins with
+    // its own kernel-ready and the session's retained notifications.
+    try {
+      await follow(session.bridge, opts.sessionGeneration);
+    } catch (error) {
+      await session.stop();
+      throw error;
+    }
+    void session.start().catch((error: unknown) => {
+      rpc.send.initializedError({ error: prettyError(error) });
     });
-    if (activeSession) {
-      (nextBridge as unknown as PyProxy).destroy();
-    } else {
+    if (!activeSession) {
       activeSession = {
         appId: opts.appId,
-        bridge: nextBridge,
+        bridge: session.bridge,
         sessionGeneration: opts.sessionGeneration,
       };
     }
@@ -107,10 +101,23 @@ async function startSession(
 }
 
 async function stopActiveSession(): Promise<void> {
-  const session = activeSession;
+  // Stopping closes every session and its notebook subscriptions.
   await self.controller.stopSession();
   activeSession = undefined;
-  (session?.bridge as unknown as PyProxy | undefined)?.destroy();
+}
+
+async function follow(
+  bridge: SerializedBridge,
+  sessionGeneration: number,
+): Promise<void> {
+  const notebook = bridge.notebook;
+  try {
+    await notebook.subscribe(`island-${sessionGeneration}`, null, (event) => {
+      rpc.send.notebookEvent({ event, sessionGeneration });
+    });
+  } finally {
+    notebook.destroy();
+  }
 }
 
 function enqueueSession<T>(operation: () => Promise<T>): Promise<T> {
@@ -243,8 +250,8 @@ export type WorkerSchema = RPCSchema<
     messages: {
       // Emitted when the worker is ready
       ready: {};
-      // Emitted when the kernel sends a message
-      kernelMessage: IslandsKernelMessage;
+      // Emitted for each event of a session's notebook stream
+      notebookEvent: IslandsNotebookEvent;
       // Emitted when the Pyodide is initialized
       initialized: {};
       // Emitted when the Pyodide fails to initialize
@@ -263,13 +270,6 @@ const rpc = createRPC<WorkerSchema, ParentSchema>({
 
 rpc.send("ready", {});
 
-/// Listeners
-// When the consumer is ready, start the message buffer
-rpc.addMessageListener("consumerReady", async () => {
-  await pyodideReadyPromise; // Make sure loading is done
-  messageBuffer.start();
-});
-
 function getMarimoVersion() {
   return self.name; // We store the version in the worker name
 }

package/src/core/kernel/__tests__/RuntimeState.test.ts

@@ -25,7 +25,7 @@ describe("RuntimeState", () => {
 
   beforeEach(() => {
     vi.clearAllMocks();
-    mockSendComponentValues = vi.fn((args) => Promise.resolve(null));
+    mockSendComponentValues = vi.fn(() => Promise.resolve(null));
     uiElementRegistry = UIElementRegistry.INSTANCE;
     uiElementRegistry.entries.clear();
 

package/src/core/kernel/__tests__/handlers.test.ts

@@ -1,7 +1,8 @@
 /* Copyright 2026 Marimo. All rights reserved. */
 import { describe, expect, it, vi } from "vitest";
 import { cellId } from "@/__tests__/branded";
-import { buildCellData, buildLayoutState } from "../handlers";
+import { createCell } from "@/core/cells/types";
+import { buildLayoutState } from "../handlers";
 import type { NotificationMessageData } from "../messages";
 import { queryParamHandlers } from "../queryParamHandlers";
 
@@ -63,152 +64,9 @@ describe("queryParamHandlers", () => {
   });
 });
 
-describe("buildCellData", () => {
-  it("should build cell data from kernel-ready data", () => {
-    const kernelReadyData: NotificationMessageData<"kernel-ready"> = {
-      cell_ids: [cellId("cell1"), cellId("cell2")],
-      codes: ["x = 1", "y = 2"],
-      names: ["__", "__"],
-      configs: [
-        { disabled: false, hide_code: false },
-        { disabled: false, hide_code: false },
-      ],
-      layout: null,
-      resumed: false,
-      ui_values: {},
-      last_executed_code: {},
-      last_execution_time: {},
-      app_config: {
-        width: "normal",
-        app_title: null,
-        layout_file: null,
-        css_file: null,
-        auto_download: [],
-      },
-      kiosk: false,
-      capabilities: {
-        terminal: false,
-      },
-      auto_instantiated: false,
-      consumer_capabilities: { edit: true, interact: true },
-    };
-
-    const cells = buildCellData(kernelReadyData);
-
-    expect(cells).toMatchInlineSnapshot(`
-      [
-        {
-          "code": "x = 1",
-          "config": {
-            "disabled": false,
-            "hide_code": false,
-          },
-          "edited": false,
-          "id": "cell1",
-          "lastCodeRun": null,
-          "lastExecutionTime": null,
-          "name": "__",
-          "serializedEditorState": null,
-        },
-        {
-          "code": "y = 2",
-          "config": {
-            "disabled": false,
-            "hide_code": false,
-          },
-          "edited": false,
-          "id": "cell2",
-          "lastCodeRun": null,
-          "lastExecutionTime": null,
-          "name": "__",
-          "serializedEditorState": null,
-        },
-      ]
-    `);
-  });
-
-  it("should mark cells as edited when code differs from last executed code", () => {
-    const kernelReadyData: NotificationMessageData<"kernel-ready"> = {
-      cell_ids: [cellId("cell1"), cellId("cell2")],
-      codes: ["x = 1", "y = 3"],
-      names: ["__", "__"],
-      configs: [
-        { disabled: false, hide_code: false },
-        { disabled: false, hide_code: false },
-      ],
-      layout: null,
-      resumed: false,
-      ui_values: {},
-      last_executed_code: {
-        cell1: "x = 1",
-        cell2: "y = 2",
-      },
-      last_execution_time: {
-        cell1: 1_234_567_890,
-        cell2: 1_234_567_891,
-      },
-      app_config: {
-        width: "normal",
-        app_title: null,
-        layout_file: null,
-        css_file: null,
-        auto_download: [],
-      },
-      kiosk: false,
-      capabilities: {
-        terminal: false,
-      },
-      auto_instantiated: false,
-      consumer_capabilities: { edit: true, interact: true },
-    };
-
-    const cells = buildCellData(kernelReadyData);
-
-    expect(cells[0].edited).toBe(false);
-    expect(cells[1].edited).toBe(true);
-    expect(cells[0].lastCodeRun).toBe("x = 1");
-    expect(cells[1].lastCodeRun).toBe("y = 2");
-  });
-
-  it("should handle empty cell data", () => {
-    const kernelReadyData: NotificationMessageData<"kernel-ready"> = {
-      cell_ids: [],
-      codes: [],
-      names: [],
-      configs: [],
-      layout: null,
-      resumed: false,
-      ui_values: {},
-      last_executed_code: {},
-      last_execution_time: {},
-      app_config: {
-        width: "normal",
-        app_title: null,
-        layout_file: null,
-        css_file: null,
-        auto_download: [],
-      },
-      kiosk: false,
-      capabilities: {
-        terminal: false,
-      },
-      auto_instantiated: false,
-      consumer_capabilities: { edit: true, interact: true },
-    };
-
-    const cells = buildCellData(kernelReadyData);
-
-    expect(cells).toMatchInlineSnapshot("[]");
-  });
-});
-
 describe("buildLayoutState", () => {
   it("should build default layout state when no layout is provided", () => {
     const kernelReadyData: NotificationMessageData<"kernel-ready"> = {
-      cell_ids: [cellId("cell1")],
-      codes: ["x = 1"],
-      names: ["__"],
-      configs: [{ disabled: false, hide_code: false }],
       layout: null,
       resumed: false,
       ui_values: {},
@@ -229,7 +87,7 @@ describe("buildLayoutState", () => {
       consumer_capabilities: { edit: true, interact: true },
     };
 
-    const cells = buildCellData(kernelReadyData);
+    const cells = [createCell({ id: cellId("cell1") })];
     const mockSetLayoutData = vi.fn();
     const layoutState = buildLayoutState(
       kernelReadyData,
@@ -248,13 +106,6 @@ describe("buildLayoutState", () => {
 
   it("should build layout state with vertical layout", () => {
     const kernelReadyData: NotificationMessageData<"kernel-ready"> = {
-      cell_ids: [cellId("cell1"), cellId("cell2")],
-      codes: ["x = 1", "y = 2"],
-      names: ["__", "__"],
-      configs: [
-        { disabled: false, hide_code: false },
-        { disabled: false, hide_code: false },
-      ],
       layout: {
         type: "vertical",
         data: [],
@@ -278,7 +129,10 @@ describe("buildLayoutState", () => {
       consumer_capabilities: { edit: true, interact: true },
     };
 
-    const cells = buildCellData(kernelReadyData);
+    const cells = [
+      createCell({ id: cellId("cell1") }),
+      createCell({ id: cellId("cell2") }),
+    ];
     const mockSetLayoutData = vi.fn();
     const layoutState = buildLayoutState(
       kernelReadyData,

package/src/core/kernel/handlers.ts

@@ -3,7 +3,8 @@ import type { LayoutType } from "@/components/editor/renderers/types";
 import { Logger } from "@/utils/Logger";
 import { Objects } from "@/utils/objects";
 import type { UIElementId } from "../cells/ids";
-import { type CellData, createCell } from "../cells/types";
+import type { ReducerCellActions } from "../cells/cells";
+import type { CellData } from "../cells/types";
 import { type AppConfig, AppConfigSchema } from "../config/config-schema";
 import { UI_ELEMENT_REGISTRY } from "../dom/uiregistry";
 import { deserializeLayoutState } from "../layout/layout";
@@ -19,41 +20,6 @@ import type { KernelState } from "./state";
 
 type KernelReadyData = NotificationMessageData<"kernel-ready">;
 
-/**
- * Build cells from kernel-ready data.
- */
-export function buildCellData(data: KernelReadyData): CellData[] {
-  const {
-    codes,
-    names,
-    configs,
-    cell_ids,
-    last_executed_code,
-    last_execution_time,
-  } = data;
-
-  const lastExecutedCode = last_executed_code || {};
-  const lastExecutionTime = last_execution_time || {};
-
-  return codes.map((code, i) => {
-    const cellId = cell_ids[i];
-    const lastCodeRun = lastExecutedCode[cellId];
-
-    // A cell is edited if the code has changed since the last time it was run
-    const edited = lastCodeRun ? lastCodeRun !== code : false;
-
-    return createCell({
-      id: cellId,
-      code,
-      edited,
-      name: names[i],
-      lastCodeRun: lastExecutedCode[cellId] ?? null,
-      lastExecutionTime: lastExecutionTime[cellId] ?? null,
-      config: configs[i],
-    });
-  });
-}
-
 /**
  * Build layout state from kernel-ready data.
  */
@@ -95,7 +61,10 @@ export function handleKernelReady(
   data: KernelReadyData,
   opts: {
     autoInstantiate: boolean;
-    setCells: (cells: CellData[], layout: LayoutState) => void;
+    /** The notebook's cells, as the host document provided them. */
+    cells: CellData[];
+    setLastRun: ReducerCellActions["setLastRun"];
+    resetCellRuntimes: ReducerCellActions["resetCellRuntimes"];
     setLayoutData: (payload: {
       layoutView: LayoutType;
       data: LayoutData;
@@ -104,17 +73,13 @@ export function handleKernelReady(
     setKernelState: (state: KernelState) => void;
     setAppConfig: (config: AppConfig) => void;
     onError: (error: Error) => void;
-    /**
-     * If provided, these cells will be used instead of the cells from
-     * kernel-ready. This allows preserving local edits made before connecting.
-     */
-    existingCells?: CellData[];
   },
-) {
+): LayoutState {
   const {
-    existingCells,
     autoInstantiate,
-    setCells,
+    cells,
+    setLastRun,
+    resetCellRuntimes,
     setLayoutData,
     setAppConfig,
     setCapabilities,
@@ -124,15 +89,12 @@ export function handleKernelReady(
   const { resumed, ui_values, app_config, capabilities, auto_instantiated } =
     data;
 
-  // Use existing cells if provided (local pre-connect edits), otherwise build from kernel-ready
-  // If the kernel was resumed, we don't want to use the existing cells because they may be stale.
-  const hasExistingCells = existingCells && existingCells.length > 0;
-  const cells =
-    hasExistingCells && !resumed ? existingCells : buildCellData(data);
-
-  // Set up layout and cells
+  resetCellRuntimes();
+  setLastRun({
+    lastCodeRun: data.last_executed_code ?? {},
+    lastExecutionTime: data.last_execution_time ?? {},
+  });
   const layoutState = buildLayoutState(data, cells, setLayoutData);
-  setCells(cells, layoutState);
 
   // Set app config and capabilities
   const parsedAppConfig = AppConfigSchema.safeParse(app_config);
@@ -145,7 +107,7 @@ export function handleKernelReady(
 
   // If the kernel was already instantiated server-side (e.g. run mode), we're done
   if (auto_instantiated) {
-    return;
+    return layoutState;
   }
 
   // If resumed, restore UI element values from kernel and we're done
@@ -156,16 +118,13 @@ export function handleKernelReady(
       const uiObjectId: UIElementId = objectId;
       UI_ELEMENT_REGISTRY.set(uiObjectId, value);
     }
-    return;
+    return layoutState;
   }
 
   // If we already have values for some objects, we should
   // send them to the kernel. This may happen after re-connecting
   // to the kernel after the computer wakes from sleep.
   const { objectIds, values } = collectUIElementValues();
-  const codesToSend = hasExistingCells
-    ? Object.fromEntries(existingCells.map((c) => [c.id, c.code]))
-    : undefined;
 
   // Send instantiate request to kernel
   void getRequestClient()
@@ -173,7 +132,6 @@ export function handleKernelReady(
       objectIds,
       values,
       autoRun: autoInstantiate,
-      codes: codesToSend,
     })
     .then(() => {
       setKernelState({ isInstantiated: true, error: null });
@@ -182,6 +140,7 @@ export function handleKernelReady(
       setKernelState({ isInstantiated: false, error: error });
       onError(new Error("Failed to instantiate", { cause: error }));
     });
+  return layoutState;
 }
 
 export function handleRemoveUIElements(

package/src/core/network/__tests__/connection.test.ts

@@ -3,7 +3,7 @@
 import { afterEach, describe, expect, it } from "vitest";
 import { initialModeAtom } from "@/core/mode";
 import { store } from "@/core/state/jotai";
-import { ConnectionState } from "../../websocket/types";
+import { ConnectionState } from "../../session/types";
 import { connectionAtom, waitForConnectionOpenIfNotebook } from "../connection";
 
 const NEVER = Symbol("never");

package/src/core/network/__tests__/requests-lazy.test.ts

@@ -53,7 +53,6 @@ describe("createLazyRequests", () => {
       sendInstantiate: vi.fn().mockResolvedValue({ instantiated: true }),
       sendFunctionRequest: vi.fn().mockResolvedValue({ result: "data" }),
       sendRestart: vi.fn().mockResolvedValue({ restarted: true }),
-      sendDeleteCell: vi.fn().mockResolvedValue({ deleted: true }),
       sendInterrupt: vi.fn().mockResolvedValue({ interrupted: true }),
       sendPdb: vi.fn().mockResolvedValue({ pdb: true }),
       sendRunScratchpad: vi.fn().mockResolvedValue({ run: true }),
@@ -320,7 +319,6 @@ describe("createLazyRequests", () => {
         cellIds: [cellId("cell1")],
         codes: ["code"],
       });
-      await lazyRequests.sendDeleteCell({ cellId: cellId("cell2") });
       await lazyRequests.sendInstantiate({
         objectIds: [uiElementId("obj1")],
         values: [],
@@ -328,7 +326,7 @@ describe("createLazyRequests", () => {
 
       // Init should only be called once
       expect(mockInit).toHaveBeenCalledTimes(1);
-      expect(mockGetRuntimeManager).toHaveBeenCalledTimes(3);
+      expect(mockGetRuntimeManager).toHaveBeenCalledTimes(2);
     });
   });
 });

package/src/core/network/connection-notice.ts

@@ -6,7 +6,7 @@ import {
   type ConnectionPhase,
   WebSocketClosedReason,
   ConnectionState,
-} from "@/core/websocket/types";
+} from "@/core/session/types";
 import { connectionAtom } from "./connection";
 
 export interface ConnectionNotice {

package/src/core/network/connection.ts

@@ -4,7 +4,7 @@ import { atomWithReducer } from "jotai/utils";
 import type { NotificationMessageData } from "../kernel/messages";
 import { isNotebookPage } from "../mode";
 import { waitFor } from "../state/jotai";
-import { type ConnectionStatus, ConnectionState } from "../websocket/types";
+import { type ConnectionStatus, ConnectionState } from "../session/types";
 
 /**
  * Shared application state, not the actual event-stream status.

package/src/core/network/requests-lazy.ts

@@ -7,7 +7,7 @@ import { memoizeLastValue } from "@/utils/once";
 import { waitForKernelToBeInstantiated } from "../kernel/state";
 import type { RuntimeManager } from "../runtime/runtime";
 import { store } from "../state/jotai";
-import { ConnectionState } from "../websocket/types";
+import { ConnectionState } from "../session/types";
 import { connectionAtom, waitForConnectionOpen } from "./connection";
 import type { EditRequests, RunRequests } from "./types";
 
@@ -52,10 +52,8 @@ const ACTIONS: Record<keyof AllRequests, Action> = {
   sendModelValue: "startConnection",
   sendInstantiate: "startConnectionWithoutKernel",
   sendRun: "startConnection",
-  sendDeleteCell: "startConnection",
   sendRunScratchpad: "startConnection",
   saveAppConfig: "startConnection",
-  saveCellConfig: "startConnection",
   discoverDataSources: "startConnection",
 
   // Export operations start a connection
@@ -144,7 +142,6 @@ const ACTIONS: Record<keyof AllRequests, Action> = {
   downloadStorage: "waitForConnectionOpen",
 
   // Sync operations that wait for connection
-  sendDocumentTransaction: "waitForConnectionOpen",
   sendCodeCompletionRequest: "waitForConnectionOpen",
 };
 

package/src/core/network/requests-network.ts

@@ -70,15 +70,6 @@ export function createNetworkRequests(): EditRequests & RunRequests {
         })
         .then(handleResponseReturnNull);
     },
-    sendDocumentTransaction: async (request) => {
-      await waitForConnectionOpen();
-      return getClient()
-        .POST("/api/document/transaction", {
-          body: request,
-          params: getParams(),
-        })
-        .then(handleResponseReturnNull);
-    },
     sendRename: (request) => {
       return getClient()
         .POST("/api/kernel/rename", {
@@ -151,14 +142,6 @@ export function createNetworkRequests(): EditRequests & RunRequests {
         })
         .then(handleResponseReturnNull);
     },
-    sendDeleteCell: (request) => {
-      return getClient()
-        .POST("/api/kernel/delete", {
-          body: request,
-          params: getParams(),
-        })
-        .then(handleResponseReturnNull);
-    },
     sendCodeCompletionRequest: async (request) => {
       await waitForConnectionOpen();
       return getClient()
@@ -184,14 +167,6 @@ export function createNetworkRequests(): EditRequests & RunRequests {
         })
         .then(handleResponseReturnNull);
     },
-    saveCellConfig: (request) => {
-      return getClient()
-        .POST("/api/kernel/set_cell_config", {
-          body: request,
-          params: getParams(),
-        })
-        .then(handleResponseReturnNull);
-    },
     sendFunctionRequest: (request) => {
       return getClient()
         .POST("/api/kernel/function_call", {

package/src/core/network/requests-static.ts

@@ -40,7 +40,6 @@ export function createStaticRequests(): EditRequests & RunRequests {
       return null;
     },
     sendRestart: throwNotInEditMode,
-    sendDocumentTransaction: throwNotInEditMode,
     sendRun: throwNotInEditMode,
     sendRunScratchpad: throwNotInEditMode,
     sendRename: throwNotInEditMode,
@@ -49,11 +48,9 @@ export function createStaticRequests(): EditRequests & RunRequests {
     sendInterrupt: throwNotInEditMode,
     sendShutdown: throwNotInEditMode,
     sendFormat: throwNotInEditMode,
-    sendDeleteCell: throwNotInEditMode,
     sendCodeCompletionRequest: throwNotInEditMode,
     saveUserConfig: throwNotInEditMode,
     saveAppConfig: throwNotInEditMode,
-    saveCellConfig: throwNotInEditMode,
     sendStdin: throwNotInEditMode,
     readCode: throwNotInEditMode,
     readSnippets: throwNotInEditMode,