@marimo-team/islands 0.25.3-dev23 → 0.25.3-dev231

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

@@ -1,10 +1,9 @@
 /* Copyright 2026 Marimo. All rights reserved. */
 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 AppConfig, AppConfigSchema } from "../config/config-schema";
+import type { ReducerCellActions } from "../cells/cells";
+import type { CellData } from "../cells/types";
 import { UI_ELEMENT_REGISTRY } from "../dom/uiregistry";
 import { deserializeLayoutState } from "../layout/layout";
 import type { LayoutData, LayoutState } from "../layout/state";
@@ -19,41 +18,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,57 +59,43 @@ 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;
     }) => void;
     setCapabilities: (capabilities: Capabilities) => void;
     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,
     setKernelState,
     onError,
   } = opts;
-  const { resumed, ui_values, app_config, capabilities, auto_instantiated } =
-    data;
+  const { resumed, ui_values, 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);
-  if (parsedAppConfig.success) {
-    setAppConfig(parsedAppConfig.data);
-  } else {
-    Logger.error("Failed to parse app config", parsedAppConfig.error);
-  }
   setCapabilities(capabilities);
 
   // 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 +106,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 +120,6 @@ export function handleKernelReady(
       objectIds,
       values,
       autoRun: autoInstantiate,
-      codes: codesToSend,
     })
     .then(() => {
       setKernelState({ isInstantiated: true, error: null });
@@ -182,6 +128,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/meta/__tests__/code-visibility.test.tsx

@@ -1,17 +1,21 @@
 /* Copyright 2026 Marimo. All rights reserved. */
 
 import { renderHook } from "@testing-library/react";
-import { createStore, Provider } from "jotai";
+import type * as Protocol from "@marimo-team/client/protocol";
+import { atom, createStore, Provider } from "jotai";
 import type { ReactNode } from "react";
 import { afterEach, beforeEach, describe, expect, it } from "vitest";
+import { partial } from "@/__mocks__/partial";
 import { showCodeInRunModeAtom } from "@/core/meta/state";
 import { type AppMode, kioskModeAtom, viewStateAtom } from "@/core/mode";
+import type { Document, NotebookReplica } from "@/core/notebook/replica";
 import {
   remoteAdapter,
   runtimeAdapterAtom,
   staticAdapter,
   type RuntimeAdapter,
 } from "@/core/runtime/adapter";
+import { activeNotebookReplicaAtom } from "@/core/session/active-notebook";
 import { useNotebookCodeAvailable } from "../code-visibility";
 
 interface StoreOpts {
@@ -19,6 +23,39 @@ interface StoreOpts {
   kiosk?: boolean;
   showInRunMode?: boolean;
   runtime?: RuntimeAdapter;
+  /** Each cell's source as the host sent it; null is withheld. */
+  codes?: ReadonlyArray<string | null>;
+}
+
+function document(codes: ReadonlyArray<string | null>): Document {
+  const cells = codes.map((code, index): Protocol.Cell => ({
+    id: `cell-${index}` as Protocol.CellId,
+    notebook_id: "nb",
+    code,
+    name: "_",
+    config: {
+      column: null,
+      disabled: false,
+      hide_code: false,
+      expand_output: false,
+    },
+    etag: `etag-${index}`,
+    code_fingerprint: `hash-${index}`,
+    version: 1,
+  }));
+  return {
+    revision: {
+      notebook_id: "nb",
+      order: cells.map((cell) => cell.id),
+      app_config: {},
+      etag: "etag",
+      saved_etag: "etag",
+      order_etag: "order",
+      app_config_etag: "config",
+      version: 1,
+    },
+    cells: new Map(cells.map((cell) => [cell.id, cell])),
+  };
 }
 
 function makeStore({
@@ -26,23 +63,32 @@ function makeStore({
   kiosk = false,
   showInRunMode = true,
   runtime = remoteAdapter,
+  codes = ["x = 1", ""],
 }: StoreOpts = {}) {
   const store = createStore();
   store.set(viewStateAtom, { mode, cellAnchor: null });
   store.set(kioskModeAtom, kiosk);
   store.set(showCodeInRunModeAtom, showInRunMode);
   store.set(runtimeAdapterAtom, runtime);
+  store.set(
+    activeNotebookReplicaAtom,
+    partial<NotebookReplica>(
+      { documentAtom: atom(document(codes)) },
+      "replica",
+    ),
+  );
   return store;
 }
 
-function wrap(store: ReturnType<typeof createStore>) {
-  return ({ children }: { children: ReactNode }) => (
+function render(store: ReturnType<typeof createStore>) {
+  const wrapper = ({ children }: { children: ReactNode }) => (
     <Provider store={store}>{children}</Provider>
   );
+  return renderHook(() => useNotebookCodeAvailable(), { wrapper }).result
+    .current;
 }
 
-const cellsWithCode = [{ code: "x = 1" }, { code: "" }];
-const cellsWithoutCode = [{ code: "" }, { code: "" }];
+const withheld = [null, null];
 
 const originalHref = window.location.href;
 
@@ -60,21 +106,11 @@ describe("useNotebookCodeAvailable", () => {
   });
 
   it("returns true in edit mode regardless of cells", () => {
-    const store = makeStore({ mode: "edit" });
-    const { result } = renderHook(
-      () => useNotebookCodeAvailable(cellsWithoutCode),
-      { wrapper: wrap(store) },
-    );
-    expect(result.current).toBe(true);
+    expect(render(makeStore({ mode: "edit", codes: withheld }))).toBe(true);
   });
 
   it("returns true in present mode (only reachable from edit)", () => {
-    const store = makeStore({ mode: "present" });
-    const { result } = renderHook(
-      () => useNotebookCodeAvailable(cellsWithoutCode),
-      { wrapper: wrap(store) },
-    );
-    expect(result.current).toBe(true);
+    expect(render(makeStore({ mode: "present", codes: withheld }))).toBe(true);
   });
 
   it("returns true in kiosk mode even when read mode and no code", () => {
@@ -82,88 +118,60 @@ describe("useNotebookCodeAvailable", () => {
       mode: "read",
       kiosk: true,
       showInRunMode: false,
+      codes: withheld,
     });
-    const { result } = renderHook(
-      () => useNotebookCodeAvailable(cellsWithoutCode),
-      { wrapper: wrap(store) },
-    );
-    expect(result.current).toBe(true);
+    expect(render(store)).toBe(true);
   });
 
   it("returns true in read mode when at least one cell has code", () => {
-    const store = makeStore({ mode: "read" });
-    const { result } = renderHook(
-      () => useNotebookCodeAvailable(cellsWithCode),
-      { wrapper: wrap(store) },
-    );
-    expect(result.current).toBe(true);
+    expect(render(makeStore({ mode: "read" }))).toBe(true);
+  });
+
+  it("counts an empty cell as code the host sent", () => {
+    expect(render(makeStore({ mode: "read", codes: ["", null] }))).toBe(true);
+  });
+
+  it("returns false in read mode when the host withheld every source", () => {
+    expect(render(makeStore({ mode: "read", codes: withheld }))).toBe(false);
   });
 
-  it("returns false in read mode when every cell.code is empty (server stripped)", () => {
+  it("returns false in read mode before a document arrives", () => {
     const store = makeStore({ mode: "read" });
-    const { result } = renderHook(
-      () => useNotebookCodeAvailable(cellsWithoutCode),
-      { wrapper: wrap(store) },
-    );
-    expect(result.current).toBe(false);
+    store.set(activeNotebookReplicaAtom, null);
+    expect(render(store)).toBe(false);
   });
 
   it("returns false in read mode when host opts out via showAppCode", () => {
-    const store = makeStore({ mode: "read", showInRunMode: false });
-    const { result } = renderHook(
-      () => useNotebookCodeAvailable(cellsWithCode),
-      { wrapper: wrap(store) },
+    expect(render(makeStore({ mode: "read", showInRunMode: false }))).toBe(
+      false,
     );
-    expect(result.current).toBe(false);
   });
 
   it("returns false in read mode when ?include-code=false", () => {
     setSearch("?include-code=false");
-    const store = makeStore({ mode: "read" });
-    const { result } = renderHook(
-      () => useNotebookCodeAvailable(cellsWithCode),
-      { wrapper: wrap(store) },
-    );
-    expect(result.current).toBe(false);
+    expect(render(makeStore({ mode: "read" }))).toBe(false);
   });
 
   it("ignores ?include-code=false in kiosk mode", () => {
     setSearch("?include-code=false");
-    const store = makeStore({ mode: "read", kiosk: true });
-    const { result } = renderHook(
-      () => useNotebookCodeAvailable(cellsWithoutCode),
-      { wrapper: wrap(store) },
-    );
-    expect(result.current).toBe(true);
+    expect(
+      render(makeStore({ mode: "read", kiosk: true, codes: withheld })),
+    ).toBe(true);
   });
 
   it("checks embedded code in a static kiosk", () => {
-    const store = makeStore({
+    const options = {
       mode: "read",
       kiosk: true,
       runtime: staticAdapter,
-    });
-
-    const withoutCode = renderHook(
-      () => useNotebookCodeAvailable(cellsWithoutCode),
-      { wrapper: wrap(store) },
-    );
-    const withCode = renderHook(() => useNotebookCodeAvailable(cellsWithCode), {
-      wrapper: wrap(store),
-    });
-
-    expect(withoutCode.result.current).toBe(false);
-    expect(withCode.result.current).toBe(true);
+    } as const;
+    expect(render(makeStore({ ...options, codes: withheld }))).toBe(false);
+    expect(render(makeStore(options))).toBe(true);
   });
 
   it("returns false for non-edit/non-read modes (home, gallery)", () => {
     for (const mode of ["home", "gallery"] as const) {
-      const store = makeStore({ mode });
-      const { result } = renderHook(
-        () => useNotebookCodeAvailable(cellsWithCode),
-        { wrapper: wrap(store) },
-      );
-      expect(result.current).toBe(false);
+      expect(render(makeStore({ mode }))).toBe(false);
     }
   });
 });

package/src/core/meta/code-visibility.ts

@@ -1,22 +1,31 @@
 /* Copyright 2026 Marimo. All rights reserved. */
 
-import { useAtomValue } from "jotai";
+import { atom, useAtomValue } from "jotai";
 import { KnownQueryParams } from "@/core/constants";
 import { showCodeInRunModeAtom } from "@/core/meta/state";
 import { kioskModeAtom, viewStateAtom } from "@/core/mode";
 import { runtimeAdapterAtom } from "@/core/runtime/adapter";
+import { confirmedDocumentAtom } from "@/core/session/active-notebook";
 import { logNever } from "@/utils/assertNever";
 
+/** Whether any cell's source reached this client. A host withholds source as null. */
+export const codeAvailableAtom = atom((get) => {
+  const document = get(confirmedDocumentAtom);
+  return (
+    document !== null &&
+    [...document.cells.values()].some((cell) => cell.code !== null)
+  );
+});
+
 /**
- * Whether the notebook source code reached the frontend and can be rendered.
+ * Whether the notebook source code can be rendered.
  *
- * In `marimo run` the server omits cell sources unless `--include-code` is
- * set (every `cell.code` arrives as `""`). Use this to gate any "show code"
- * / "copy code" / "download .py" affordance.
+ * In `marimo run` the server withholds cell sources unless `--include-code`
+ * is set. Use this to gate any "show code" / "copy code" / "download .py"
+ * affordance.
  */
-export function useNotebookCodeAvailable(
-  cells: ReadonlyArray<{ code: string }>,
-): boolean {
+export function useNotebookCodeAvailable(): boolean {
+  const available = useAtomValue(codeAvailableAtom);
   const kioskMode = useAtomValue(kioskModeAtom);
   const { mode } = useAtomValue(viewStateAtom);
   const showInRunMode = useAtomValue(showCodeInRunModeAtom);
@@ -41,7 +50,7 @@ export function useNotebookCodeAvailable(
       if (params.get(KnownQueryParams.includeCode) === "false") {
         return false;
       }
-      return cells.some((cell) => Boolean(cell.code));
+      return available;
     }
     default:
       logNever(mode);

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 {