@marimo-team/islands 0.23.15-dev4 → 0.23.15-dev40

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

@@ -0,0 +1,214 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+import { act } from "react";
+import {
+  afterAll,
+  afterEach,
+  beforeAll,
+  beforeEach,
+  describe,
+  expect,
+  it,
+} from "vitest";
+import { cellId } from "@/__tests__/branded";
+import { MockNotebook } from "@/__mocks__/notebook";
+import { MockRequestClient } from "@/__mocks__/requests";
+import { notebookAtom } from "@/core/cells/cells";
+import {
+  ISLAND_DATA_ATTRIBUTES,
+  ISLAND_TAG_NAMES,
+  ISLANDS_JSON_SCRIPT_TYPE,
+} from "@/core/islands/constants";
+import { requestClientAtom } from "@/core/network/requests";
+import { store } from "@/core/state/jotai";
+import { parseMarimoIslandApps } from "../../parse";
+import { MarimoIslandElement } from "../web-components";
+
+beforeAll(() => {
+  setReactActEnvironment(true);
+  if (!customElements.get(ISLAND_TAG_NAMES.ISLAND)) {
+    customElements.define(ISLAND_TAG_NAMES.ISLAND, MarimoIslandElement);
+  }
+});
+
+afterAll(() => {
+  setReactActEnvironment(false);
+});
+
+beforeEach(() => {
+  store.set(notebookAtom, MockNotebook.notebookState({ cellData: {} }));
+  store.set(requestClientAtom, MockRequestClient.create());
+});
+
+afterEach(async () => {
+  await actAndFlush(() => document.body.replaceChildren());
+  store.set(requestClientAtom, null);
+});
+
+describe("MarimoIslandElement lifecycle", () => {
+  it("binds a reactive island after its runtime cell index is materialized", async () => {
+    setNotebook("outgoing-cell", "outgoing runtime output");
+    const container = document.createElement("div");
+    container.innerHTML = `
+      <marimo-island data-app-id="app-1" data-reactive="true">
+        <marimo-cell-output><div>initial output</div></marimo-cell-output>
+        <marimo-cell-code hidden>${encodeURIComponent("value = 1")}</marimo-cell-code>
+      </marimo-island>
+    `;
+    await actAndFlush(() => document.body.append(container));
+
+    const island = container.querySelector<HTMLElement>(
+      ISLAND_TAG_NAMES.ISLAND,
+    );
+    expect(island).not.toBeNull();
+    expect(island?.getAttribute("data-status")).toBeNull();
+
+    await actAndFlush(() => {
+      parseMarimoIslandApps(container);
+    });
+
+    expect(island?.textContent).toContain("outgoing runtime output");
+
+    await actAndFlush(() => {
+      setNotebook("runtime-cell");
+    });
+
+    expect(island?.getAttribute(ISLAND_DATA_ATTRIBUTES.CELL_IDX)).toBe("0");
+    expect(island?.getAttribute("data-status")).toBe("idle");
+
+    await actAndFlush(() => {
+      island?.remove();
+      if (island) {
+        container.append(island);
+      }
+    });
+    expect(island?.textContent).toContain("initial output");
+
+    island?.setAttribute(ISLAND_DATA_ATTRIBUTES.CELL_ID, "source-cell");
+    island?.setAttribute(ISLAND_DATA_ATTRIBUTES.REACTIVE, "false");
+    expect((island as MarimoIslandElement | null)?.cellId).toBeUndefined();
+  });
+
+  it("refreshes a reused payload island and clears an old app binding", async () => {
+    const container = document.createElement("div");
+    container.innerHTML = `
+      <marimo-island
+        data-app-id="app-1"
+        data-cell-id="cell-1"
+        data-reactive="true"
+      >
+        <marimo-cell-output><div>old output</div></marimo-cell-output>
+        <marimo-cell-code hidden>${encodeURIComponent('print("old")')}</marimo-cell-code>
+      </marimo-island>
+    `;
+    const script = document.createElement("script");
+    script.type = ISLANDS_JSON_SCRIPT_TYPE;
+    script.textContent = payloadSource({
+      code: 'print("unchanged")',
+      outputHtml: "<div>old output</div>",
+    });
+    container.append(script);
+    await actAndFlush(() => document.body.append(container));
+
+    await actAndFlush(() => {
+      parseMarimoIslandApps(container);
+      script.textContent = payloadSource({
+        code: 'print("unchanged")',
+        outputHtml: "<div>updated output</div>",
+      });
+      parseMarimoIslandApps(container);
+    });
+
+    const island = container.querySelector<HTMLElement>(
+      ISLAND_TAG_NAMES.ISLAND,
+    );
+    expect(island?.textContent).toContain("updated output");
+    expect(island?.textContent).not.toContain("old output");
+    expect(island?.querySelector(ISLAND_TAG_NAMES.CELL_OUTPUT)).toBeNull();
+
+    await actAndFlush(() => {
+      setNotebook("outgoing-cell", "outgoing runtime output");
+    });
+    expect(island?.textContent).toContain("outgoing runtime output");
+
+    await actAndFlush(() => {
+      island?.setAttribute(ISLAND_DATA_ATTRIBUTES.APP_ID, "app-2");
+      island?.setAttribute(ISLAND_DATA_ATTRIBUTES.CELL_ID, "cell-2");
+      script.textContent = payloadSource({
+        appId: "app-2",
+        cellId: "cell-2",
+        code: 'print("next")',
+        outputHtml: "<div>next app output</div>",
+      });
+      parseMarimoIslandApps(container);
+    });
+
+    expect(island?.textContent).toContain("next app output");
+    expect(island?.textContent).not.toContain("outgoing runtime output");
+  });
+});
+
+function setNotebook(id: string, output?: string): void {
+  const runtimeCellId = cellId(id);
+  store.set(
+    notebookAtom,
+    MockNotebook.notebookState({
+      cellData: { [runtimeCellId]: {} },
+      cellRuntime: output
+        ? {
+            [runtimeCellId]: {
+              output: {
+                channel: "output",
+                data: `<div>${output}</div>`,
+                mimetype: "text/html",
+                timestamp: 0,
+              },
+            },
+          }
+        : undefined,
+    }),
+  );
+}
+
+function payloadSource({
+  appId = "app-1",
+  cellId = "cell-1",
+  code,
+  outputHtml,
+}: {
+  appId?: string;
+  cellId?: string;
+  code: string;
+  outputHtml: string;
+}): string {
+  return JSON.stringify({
+    schemaVersion: 1,
+    appId,
+    cells: [
+      {
+        cellId,
+        code,
+        outputHtml,
+        outputMimetype: "text/html",
+        reactive: true,
+        displayCode: false,
+        displayOutput: true,
+      },
+    ],
+  });
+}
+
+async function actAndFlush(action: () => void): Promise<void> {
+  await act(async () => {
+    action();
+    await Promise.resolve();
+  });
+}
+
+function setReactActEnvironment(value: boolean): void {
+  (
+    globalThis as typeof globalThis & {
+      IS_REACT_ACT_ENVIRONMENT: boolean;
+    }
+  ).IS_REACT_ACT_ENVIRONMENT = value;
+}

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

@@ -5,7 +5,7 @@ import { isValidElement, type JSX } from "react";
 import ReactDOM, { type Root } from "react-dom/client";
 import { ErrorBoundary } from "@/components/editor/boundary/ErrorBoundary";
 import { TooltipProvider } from "@/components/ui/tooltip";
-import { notebookAtom } from "@/core/cells/cells";
+import { cellIdsAtom, notebookAtom } from "@/core/cells/cells";
 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";
@@ -16,9 +16,10 @@ import { store } from "../../state/jotai";
 import {
   ISLAND_CSS_CLASSES,
   ISLAND_DATA_ATTRIBUTES,
+  ISLAND_SOURCE_CHANGED_EVENT,
   ISLAND_TAG_NAMES,
 } from "../constants";
-import { extractIslandCodeFromEmbed } from "../parse";
+import { extractIslandCodeFromEmbed, retainIslandSource } from "../parse";
 import { MarimoOutputWrapper } from "./output-wrapper";
 
 /**
@@ -28,7 +29,6 @@ export interface IslandRenderConfig {
   html: string;
   codeCallback: () => string;
   editor: JSX.Element | null;
-  cellId: CellId | undefined;
 }
 
 /**
@@ -38,7 +38,12 @@ export interface IslandRenderConfig {
  * functionality like re-running cells and copying code.
  */
 export class MarimoIslandElement extends HTMLElement {
+  private connectionGeneration = 0;
   private root?: Root;
+  private renderConfig?: IslandRenderConfig;
+  private runtimeAppId?: string;
+  private runtimeCellId?: CellId;
+  private unsubscribeCellIds?: () => void;
 
   public static readonly tagName = ISLAND_TAG_NAMES.ISLAND;
   public static readonly outputTagName = ISLAND_TAG_NAMES.CELL_OUTPUT;
@@ -72,6 +77,9 @@ export class MarimoIslandElement extends HTMLElement {
    * Returns undefined for non-reactive islands (they have no corresponding cell).
    */
   get cellId(): CellId | undefined {
+    if (!this.isReactive) {
+      return undefined;
+    }
     const cellId = this.getAttribute(ISLAND_DATA_ATTRIBUTES.CELL_ID);
     if (cellId) {
       return cellId as CellId;
@@ -94,11 +102,9 @@ export class MarimoIslandElement extends HTMLElement {
   /**
    * Looks up a cell ID from the notebook state by index
    */
-  private getCellIdFromIndex(idx: number): CellId {
+  private getCellIdFromIndex(idx: number): CellId | undefined {
     const { cellIds } = store.get(notebookAtom);
-    const cellId = cellIds.inOrderIds.at(idx);
-    invariant(cellId, `Missing cell ID at index ${idx}`);
-    return cellId;
+    return cellIds.inOrderIds.at(idx);
   }
 
   /**
@@ -110,18 +116,62 @@ export class MarimoIslandElement extends HTMLElement {
    * synchronously from there causes "unmount during render" warnings.
    */
   connectedCallback(): void {
-    // Capture config synchronously (before children get cleared by createRoot)
-    const config = this.extractRenderConfig();
+    const connectionGeneration = ++this.connectionGeneration;
+    this.addEventListener(
+      ISLAND_SOURCE_CHANGED_EVENT,
+      this.handleSourceChanged,
+    );
+    // Capture config synchronously (before children get cleared by createRoot).
+    // A reconnected element reuses the source captured on its first mount.
+    if (this.querySelector(MarimoIslandElement.outputTagName)) {
+      this.renderConfig = this.extractRenderConfig();
+    }
+    invariant(this.renderConfig, "Missing island render source");
+    this.runtimeAppId = this.appId;
+    this.runtimeCellId = this.hasAttribute(ISLAND_DATA_ATTRIBUTES.CELL_IDX)
+      ? this.cellId
+      : undefined;
+    this.syncCellIdsSubscription();
     queueMicrotask(() => {
       // Guard against disconnect between connectedCallback and microtask
-      if (!this.isConnected) {
+      if (
+        !this.isConnected ||
+        connectionGeneration !== this.connectionGeneration
+      ) {
         return;
       }
       this.root = ReactDOM.createRoot(this);
-      this.renderIsland(config);
+      this.renderIsland();
     });
   }
 
+  private handleSourceChanged = (): void => {
+    if (this.querySelector(MarimoIslandElement.outputTagName)) {
+      this.renderConfig = this.extractRenderConfig();
+      this.querySelector(MarimoIslandElement.outputTagName)?.remove();
+      this.querySelector(MarimoIslandElement.codeTagName)?.remove();
+    }
+    this.runtimeCellId =
+      this.isReactive && this.runtimeAppId === this.appId
+        ? this.cellId
+        : undefined;
+    this.runtimeAppId = this.appId;
+    this.syncCellIdsSubscription();
+    this.renderIsland();
+  };
+
+  private syncCellIdsSubscription(): void {
+    if (this.isReactive) {
+      this.unsubscribeCellIds ??= store.sub(cellIdsAtom, () => {
+        this.runtimeCellId = this.cellId;
+        this.renderIsland();
+      });
+      return;
+    }
+    this.unsubscribeCellIds?.();
+    this.unsubscribeCellIds = undefined;
+  }
+
   /**
    * Extracts configuration needed for rendering
    */
@@ -130,7 +180,7 @@ export class MarimoIslandElement extends HTMLElement {
     const initialOutput = output.innerHTML;
     const optionalEditor = this.getOptionalEditor();
     const code = this.code;
-    const cellId = this.cellId;
+    retainIslandSource(this, { code, output: initialOutput });
 
     // Read objectId directly from the DOM before createRoot clears children.
     // optionalEditor is a <RenderHTML> wrapper, so its .props don't carry the
@@ -152,15 +202,19 @@ export class MarimoIslandElement extends HTMLElement {
       html: initialOutput,
       codeCallback,
       editor: optionalEditor,
-      cellId,
     };
   }
 
   /**
    * Renders the island with React
    */
-  private renderIsland(config: IslandRenderConfig): void {
-    const { html, codeCallback, editor, cellId } = config;
+  private renderIsland(): void {
+    const config = this.renderConfig;
+    if (!config) {
+      return;
+    }
+    const { html, codeCallback, editor } = config;
+    const cellId = this.runtimeCellId;
     const alwaysShowRun = !!editor;
     const trimmedHtml = html.trim();
     const isEmpty = trimmedHtml === "<span></span>" || trimmedHtml === "";
@@ -168,6 +222,7 @@ export class MarimoIslandElement extends HTMLElement {
 
     // Non-reactive islands have no cell in the kernel — just render static HTML
     if (!cellId) {
+      this.removeAttribute("data-status");
       this.root?.render(
         <ErrorBoundary>
           <Provider store={store}>
@@ -242,6 +297,12 @@ export class MarimoIslandElement extends HTMLElement {
    * Cleanup when element is removed from DOM
    */
   disconnectedCallback(): void {
+    this.removeEventListener(
+      ISLAND_SOURCE_CHANGED_EVENT,
+      this.handleSourceChanged,
+    );
+    this.unsubscribeCellIds?.();
+    this.unsubscribeCellIds = undefined;
     const root = this.root;
     this.root = undefined;
     // Defer unmount to avoid "unmount during render" race

package/src/core/islands/constants.ts

@@ -11,6 +11,7 @@ export const ISLAND_TAG_NAMES = {
 } as const;
 
 export const ISLANDS_JSON_SCRIPT_TYPE = "application/vnd.marimo.islands+json";
+export const ISLAND_SOURCE_CHANGED_EVENT = "marimo-island-source-changed";
 
 /**
  * Data attributes for islands

package/src/core/islands/main.ts

@@ -14,12 +14,78 @@ import "iconify-icon";
 import { Logger } from "@/utils/Logger";
 import { initializeIslands } from "./bootstrap";
 import { getGlobalBridge } from "./bridge";
+import { ISLAND_CSS_CLASSES, ISLAND_TAG_NAMES } from "./constants";
+import { parseMarimoIslandApps } from "./parse";
+
+const bridge = getGlobalBridge();
+let bootstrapPromise: Promise<void> | undefined;
+
+/**
+ * Public entry point for the islands bundle.
+ *
+ * Islands auto-initialize on load. Hosts that keep the browser realm alive
+ * across client-side navigation (retaining this worker and its Pyodide
+ * environment) can drive page transitions with:
+ *
+ * ```ts
+ * if (canReplaceApp()) {
+ *   await stopApp();
+ *   updatePage();
+ *   await initialize();
+ * }
+ * ```
+ *
+ * The retained lifecycle (`canReplaceApp`/`stopApp`) only applies to pages
+ * whose islands resolve to a single app; multi-app pages are not replaceable.
+ */
+
+/**
+ * Returns whether the current document can replace its app in this worker.
+ *
+ * True only when the document contains exactly one replaceable app. This is a
+ * read-only probe: it does not materialize island payloads or mutate the DOM.
+ */
+export function canReplaceApp(): boolean {
+  if (!document.querySelector(ISLAND_TAG_NAMES.ISLAND)) {
+    return false;
+  }
+  return parseMarimoIslandApps(document, { materialize: false }).length === 1;
+}
+
+/**
+ * Mounts marimo custom elements and starts the apps in the current document.
+ *
+ * Safe to call repeatedly, including after client-side navigation: the worker
+ * bootstrap is memoized (and cleared if it fails, so a later call retries),
+ * while app discovery re-runs on every call to pick up the current DOM.
+ */
+export async function initialize(): Promise<void> {
+  const islands = document.querySelectorAll<HTMLElement>(
+    ISLAND_TAG_NAMES.ISLAND,
+  );
+  if (islands.length === 0) {
+    return;
+  }
+  for (const island of islands) {
+    island.classList.add(ISLAND_CSS_CLASSES.NAMESPACE);
+  }
+  bootstrapPromise ??= initializeIslands({ bridge }).catch((error: unknown) => {
+    bootstrapPromise = undefined;
+    throw error;
+  });
+  await bootstrapPromise;
+  await bridge.initializeApps();
+}
 
 /**
- * Main entry point for the js bundle for embedded marimo apps.
+ * Stops the matching active app session.
+ *
+ * With no `appId`, stops the current retained single-app session; with an
+ * `appId`, stops it only if it matches the active app. No-op when there is no
+ * retained session to stop.
  */
-export async function initialize() {
-  await initializeIslands({ bridge: getGlobalBridge() });
+export async function stopApp(appId?: string): Promise<void> {
+  await bridge.stopSession(appId);
 }
 
 // Auto-initialize on module load

package/src/core/islands/parse.ts

@@ -2,6 +2,7 @@
 
 import {
   ISLAND_DATA_ATTRIBUTES,
+  ISLAND_SOURCE_CHANGED_EVENT,
   ISLAND_TAG_NAMES,
   ISLANDS_JSON_SCRIPT_TYPE,
 } from "@/core/islands/constants";
@@ -74,12 +75,15 @@ interface MarimoIslandPayloadCell {
   displayOutput: boolean;
 }
 
+const retainedPayloadCellIds = new WeakMap<HTMLElement, string>();
+
 /**
  * Parses marimo island apps from the DOM
  * @param root - Root element to search within (defaults to document)
  */
 export function parseMarimoIslandApps(
   root: Document | Element = document,
+  options: { materialize?: boolean } = {},
 ): MarimoIslandApp[] {
   const embeds = [
     ...root.querySelectorAll<HTMLElement>(ISLAND_TAG_NAMES.ISLAND),
@@ -90,11 +94,12 @@ export function parseMarimoIslandApps(
     return [];
   }
 
+  const materialize = options.materialize ?? true;
   if (payloads.length > 0) {
-    return parsePayloadBackedApps(embeds, payloads);
+    return parsePayloadBackedApps({ embeds, materialize, payloads });
   }
 
-  return parseIslandElementsIntoApps(embeds);
+  return parseIslandElementsIntoApps(embeds, materialize);
 }
 
 /**
@@ -103,6 +108,7 @@ export function parseMarimoIslandApps(
  */
 export function parseIslandElementsIntoApps(
   embeds: HTMLElement[],
+  materialize = true,
 ): MarimoIslandApp[] {
   const apps = new Map<string, MarimoIslandApp>();
 
@@ -140,16 +146,24 @@ export function parseIslandElementsIntoApps(
     });
 
     // Add data-cell-idx attribute to the island element
-    embed.setAttribute(ISLAND_DATA_ATTRIBUTES.CELL_IDX, idx.toString());
+    if (materialize) {
+      embed.setAttribute(ISLAND_DATA_ATTRIBUTES.CELL_IDX, idx.toString());
+      embed.dispatchEvent(new Event(ISLAND_SOURCE_CHANGED_EVENT));
+    }
   }
 
   return [...apps.values()];
 }
 
-function parsePayloadBackedApps(
-  embeds: HTMLElement[],
-  payloads: MarimoIslandPayload[],
-): MarimoIslandApp[] {
+function parsePayloadBackedApps({
+  embeds,
+  materialize,
+  payloads,
+}: {
+  embeds: HTMLElement[];
+  materialize: boolean;
+  payloads: MarimoIslandPayload[];
+}): MarimoIslandApp[] {
   const apps = new Map<string, MarimoIslandApp>();
   const matchedPayloadCells = new Map<MarimoIslandPayloadCell, HTMLElement>();
   const consumedEmbeds = new Set<HTMLElement>();
@@ -169,7 +183,9 @@ function parsePayloadBackedApps(
       }
       consumedEmbeds.add(embed);
       matchedPayloadCells.set(cell, embed);
-      materializeIslandPayload(embed, cell);
+      if (materialize) {
+        materializeIslandPayload(embed, cell);
+      }
       hasMatchedIsland = true;
     }
     // Only payloads matched to island anchors can start runtime apps.
@@ -215,7 +231,7 @@ function parsePayloadBackedApps(
         appCell.disabled = true;
       }
       app.cells.push(appCell);
-      if (cell.reactive) {
+      if (materialize && cell.reactive) {
         embed?.setAttribute(ISLAND_DATA_ATTRIBUTES.CELL_IDX, idx.toString());
       }
     }
@@ -223,12 +239,15 @@ function parsePayloadBackedApps(
 
   // A supported payload is the runtime source for its app. Extra same-app DOM
   // islands are disconnected from runtime binding.
-  for (const embed of embeds) {
-    const appId = embed.getAttribute(ISLAND_DATA_ATTRIBUTES.APP_ID);
-    if (appId && payloadAppIds.has(appId) && !consumedEmbeds.has(embed)) {
-      embed.removeAttribute(ISLAND_DATA_ATTRIBUTES.CELL_ID);
-      embed.removeAttribute(ISLAND_DATA_ATTRIBUTES.CELL_IDX);
-      embed.setAttribute(ISLAND_DATA_ATTRIBUTES.REACTIVE, "false");
+  if (materialize) {
+    for (const embed of embeds) {
+      const appId = embed.getAttribute(ISLAND_DATA_ATTRIBUTES.APP_ID);
+      if (appId && payloadAppIds.has(appId) && !consumedEmbeds.has(embed)) {
+        embed.removeAttribute(ISLAND_DATA_ATTRIBUTES.CELL_ID);
+        embed.removeAttribute(ISLAND_DATA_ATTRIBUTES.CELL_IDX);
+        embed.setAttribute(ISLAND_DATA_ATTRIBUTES.REACTIVE, "false");
+        consumedEmbeds.add(embed);
+      }
     }
   }
 
@@ -237,7 +256,16 @@ function parsePayloadBackedApps(
     return !appId || !payloadAppIds.has(appId);
   });
 
-  return [...apps.values(), ...parseIslandElementsIntoApps(domOnlyEmbeds)];
+  if (materialize) {
+    for (const embed of consumedEmbeds) {
+      embed.dispatchEvent(new Event(ISLAND_SOURCE_CHANGED_EVENT));
+    }
+  }
+
+  return [
+    ...apps.values(),
+    ...parseIslandElementsIntoApps(domOnlyEmbeds, materialize),
+  ];
 }
 
 function findMatchingIsland({
@@ -257,7 +285,8 @@ function findMatchingIsland({
     }
     return (
       embed.getAttribute(ISLAND_DATA_ATTRIBUTES.APP_ID) === appId &&
-      embed.getAttribute(ISLAND_DATA_ATTRIBUTES.CELL_ID) === cell.cellId
+      (embed.getAttribute(ISLAND_DATA_ATTRIBUTES.CELL_ID) ??
+        retainedPayloadCellIds.get(embed)) === cell.cellId
     );
   });
 }
@@ -266,6 +295,7 @@ function materializeIslandPayload(
   embed: HTMLElement,
   cell: MarimoIslandPayloadCell,
 ): void {
+  retainedPayloadCellIds.set(embed, cell.cellId);
   embed.setAttribute(
     ISLAND_DATA_ATTRIBUTES.REACTIVE,
     JSON.stringify(cell.reactive),
@@ -304,6 +334,22 @@ function ensureIslandChild(embed: HTMLElement, tagName: string): HTMLElement {
  * @param embed - The island HTML element
  * @returns Cell data or null if invalid
  */
+const retainedIslandSources = new WeakMap<
+  HTMLElement,
+  { output: string; code: string }
+>();
+
+export function retainIslandSource(
+  embed: HTMLElement,
+  source: { output: string; code: string },
+): void {
+  if (source.code) {
+    retainedIslandSources.set(embed, source);
+  } else {
+    retainedIslandSources.delete(embed);
+  }
+}
+
 export function parseIslandElement(
   embed: HTMLElement,
 ): { output: string; code: string } | null {
@@ -312,14 +358,15 @@ export function parseIslandElement(
   );
   const code = extractIslandCodeFromEmbed(embed);
 
-  if (!cellOutput || !code) {
-    return null;
+  if (cellOutput && code) {
+    return {
+      output: cellOutput.innerHTML,
+      code: code,
+    };
   }
 
-  return {
-    output: cellOutput.innerHTML,
-    code: code,
-  };
+  const retained = retainedIslandSources.get(embed);
+  return retained?.code ? retained : null;
 }
 
 export function createMarimoFile(app: {