From 4faa829c40e82c17d227d38042282dd78e612184 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?P=C3=A9ter=20Ferenc=20Gyarmati?= Date: Sat, 11 Jul 2026 20:32:32 +0200 Subject: [PATCH 01/12] refactor(islands): make app discovery repeatable Allow callers to inspect island apps without materializing payloads. Retain payload identities and rendered source so repeated initialization reconstructs the same app after custom elements take ownership of their children. --- .../src/core/islands/__tests__/parse.test.ts | 64 +++++++++++++++++ .../islands/components/web-components.tsx | 3 +- frontend/src/core/islands/parse.ts | 72 ++++++++++++++----- 3 files changed, 121 insertions(+), 18 deletions(-) diff --git a/frontend/src/core/islands/__tests__/parse.test.ts b/frontend/src/core/islands/__tests__/parse.test.ts index 593170e28fb..547f429b2f5 100644 --- a/frontend/src/core/islands/__tests__/parse.test.ts +++ b/frontend/src/core/islands/__tests__/parse.test.ts @@ -15,6 +15,7 @@ import { parseIslandElement, parseIslandElementsIntoApps, parseMarimoIslandApps, + retainIslandSource, } from "../parse"; import { createMockIslandElement, createMockIslands } from "./test-utils.tsx"; @@ -317,6 +318,19 @@ describe("parseIslandElement", () => { expect(result).toBeNull(); }); + + it("uses retained source after the custom element renders", () => { + const element = document.createElement(ISLAND_TAG_NAMES.ISLAND); + retainIslandSource(element, { + code: 'print("retained")', + output: "
retained output
", + }); + + expect(parseIslandElement(element)).toEqual({ + code: 'print("retained")', + output: "
retained output
", + }); + }); }); describe("parseIslandElementsIntoApps", () => { @@ -630,6 +644,56 @@ describe("parseMarimoIslandApps", () => { ); }); + it("preserves payload-backed cells after a non-materializing capability probe", () => { + const island = createMockIslandElement({ + appId: "app1", + cellId: "cell-2", + code: "dom_code = True", + innerHTML: "
dom output
", + }); + island.setAttribute(ISLAND_DATA_ATTRIBUTES.REACTIVE, "true"); + container.appendChild(island); + appendPayload(container, { + schemaVersion: 1, + appId: "app1", + cells: [createPayloadCell({ cellId: "cell-2" })], + }); + const originalIsland = island.outerHTML; + + const probed = parseMarimoIslandApps(container, { materialize: false }); + + expect(probed).toHaveLength(1); + expect(island.outerHTML).toBe(originalIsland); + + expect(parseMarimoIslandApps(container)).toEqual(probed); + expect(island.getAttribute(ISLAND_DATA_ATTRIBUTES.CELL_ID)).toBeNull(); + expect(island.getAttribute(ISLAND_DATA_ATTRIBUTES.CELL_IDX)).toBe("0"); + }); + + it("keeps a payload anchor stable during duplicate initialization", () => { + const island = createMockIslandElement({ + appId: "app1", + cellId: "cell-1", + code: "dom_code = True", + innerHTML: "
dom output
", + }); + island.setAttribute(ISLAND_DATA_ATTRIBUTES.REACTIVE, "true"); + container.appendChild(island); + appendPayload(container, { + schemaVersion: 1, + appId: "app1", + cells: [createPayloadCell()], + }); + + const first = parseMarimoIslandApps(container); + island.replaceChildren(); + const second = parseMarimoIslandApps(container); + + expect(second).toEqual(first); + expect(island.childNodes).toHaveLength(0); + expect(island.getAttribute(ISLAND_DATA_ATTRIBUTES.CELL_IDX)).toBe("0"); + }); + it("should parse Python-generated island payload snapshots", () => { const html = readFileSync( new URL( diff --git a/frontend/src/core/islands/components/web-components.tsx b/frontend/src/core/islands/components/web-components.tsx index 557b5e90027..88a9d500fca 100644 --- a/frontend/src/core/islands/components/web-components.tsx +++ b/frontend/src/core/islands/components/web-components.tsx @@ -18,7 +18,7 @@ import { ISLAND_DATA_ATTRIBUTES, ISLAND_TAG_NAMES, } from "../constants"; -import { extractIslandCodeFromEmbed } from "../parse"; +import { extractIslandCodeFromEmbed, retainIslandSource } from "../parse"; import { MarimoOutputWrapper } from "./output-wrapper"; /** @@ -131,6 +131,7 @@ export class MarimoIslandElement extends HTMLElement { 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 wrapper, so its .props don't carry the diff --git a/frontend/src/core/islands/parse.ts b/frontend/src/core/islands/parse.ts index 6662b70f665..a0f83725d1c 100644 --- a/frontend/src/core/islands/parse.ts +++ b/frontend/src/core/islands/parse.ts @@ -74,12 +74,15 @@ interface MarimoIslandPayloadCell { displayOutput: boolean; } +const retainedPayloadCellIds = new WeakMap(); + /** * 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(ISLAND_TAG_NAMES.ISLAND), @@ -90,11 +93,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 +107,7 @@ export function parseMarimoIslandApps( */ export function parseIslandElementsIntoApps( embeds: HTMLElement[], + materialize = true, ): MarimoIslandApp[] { const apps = new Map(); @@ -140,16 +145,23 @@ 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()); + } } 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(); const matchedPayloadCells = new Map(); const consumedEmbeds = new Set(); @@ -169,7 +181,9 @@ function parsePayloadBackedApps( } consumedEmbeds.add(embed); matchedPayloadCells.set(cell, embed); - materializeIslandPayload(embed, cell); + if (materialize && !retainedPayloadCellIds.has(embed)) { + materializeIslandPayload(embed, cell); + } hasMatchedIsland = true; } // Only payloads matched to island anchors can start runtime apps. @@ -215,7 +229,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 +237,14 @@ 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"); + } } } @@ -237,7 +253,10 @@ function parsePayloadBackedApps( return !appId || !payloadAppIds.has(appId); }); - return [...apps.values(), ...parseIslandElementsIntoApps(domOnlyEmbeds)]; + return [ + ...apps.values(), + ...parseIslandElementsIntoApps(domOnlyEmbeds, materialize), + ]; } function findMatchingIsland({ @@ -257,7 +276,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 +286,7 @@ function materializeIslandPayload( embed: HTMLElement, cell: MarimoIslandPayloadCell, ): void { + retainedPayloadCellIds.set(embed, cell.cellId); embed.setAttribute( ISLAND_DATA_ATTRIBUTES.REACTIVE, JSON.stringify(cell.reactive), @@ -304,9 +325,26 @@ 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 { + retainedIslandSources.set(embed, source); +} + export function parseIslandElement( embed: HTMLElement, ): { output: string; code: string } | null { + const retained = retainedIslandSources.get(embed); + if (retained) { + return retained.code ? retained : null; + } + const cellOutput = embed.querySelector( ISLAND_TAG_NAMES.CELL_OUTPUT, ); From 4aa58ab6930dc052e7939c2f510f60e4231a1f71 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?P=C3=A9ter=20Ferenc=20Gyarmati?= Date: Sat, 11 Jul 2026 20:32:45 +0200 Subject: [PATCH 02/12] feat(wasm): add session teardown for retained runtimes Own the Python session task and stop callable so an app can be cancelled and collected while Pyodide remains active. Serialize dependency loading and skip initialization work from superseded sessions. --- .../worker/__tests__/controller.test.ts | 96 ++++++++++++++ frontend/src/core/wasm/worker/bootstrap.ts | 119 +++++++++++++++--- 2 files changed, 195 insertions(+), 20 deletions(-) create mode 100644 frontend/src/core/islands/worker/__tests__/controller.test.ts diff --git a/frontend/src/core/islands/worker/__tests__/controller.test.ts b/frontend/src/core/islands/worker/__tests__/controller.test.ts new file mode 100644 index 00000000000..a583a292443 --- /dev/null +++ b/frontend/src/core/islands/worker/__tests__/controller.test.ts @@ -0,0 +1,96 @@ +/* Copyright 2026 Marimo. All rights reserved. */ + +import type { PyodideInterface } from "pyodide"; +import { describe, expect, it, vi } from "vitest"; +import { ReadonlyWasmController } from "../controller"; + +class TestController extends ReadonlyWasmController { + setPyodide(pyodide: PyodideInterface) { + this.pyodide = pyodide; + } +} + +function createPyodideStub() { + const init = Object.assign(vi.fn(), { destroy: vi.fn() }); + const stop = Object.assign(vi.fn().mockResolvedValue(undefined), { + destroy: vi.fn(), + }); + const packages = { destroy: vi.fn(), toJs: () => [] }; + const sessionResources = Object.assign([{}, init, packages, stop], { + destroy: vi.fn(), + }); + const loadPackagesFromImports = vi.fn().mockResolvedValue(undefined); + const pyodide = { + runPython: vi.fn(() => sessionResources), + loadPackagesFromImports, + loadedPackages: {}, + runPythonAsync: vi.fn(), + } as unknown as PyodideInterface; + + return { + init, + loadPackagesFromImports, + packages, + pyodide, + sessionResources, + stop, + }; +} + +function startSession(controller: TestController, code: string) { + return controller.startSession({ + code, + filename: `${code}.py`, + onMessage: vi.fn(), + }); +} + +describe("WASM controller session lifecycle", () => { + it("stops and releases the active session", async () => { + const { init, packages, pyodide, sessionResources, stop } = + createPyodideStub(); + const controller = new TestController(); + controller.setPyodide(pyodide); + + await startSession(controller, "current"); + await vi.waitFor(() => expect(init).toHaveBeenCalledOnce()); + await controller.stopSession(); + + expect(stop).toHaveBeenCalledOnce(); + expect(stop.destroy).toHaveBeenCalledOnce(); + expect(sessionResources.destroy).toHaveBeenCalledOnce(); + expect(packages.destroy).toHaveBeenCalledOnce(); + expect(init.destroy).toHaveBeenCalledOnce(); + }); + + it("loads dependencies serially and skips superseded sessions", async () => { + let finishFirstLoad!: () => void; + const { loadPackagesFromImports, pyodide } = createPyodideStub(); + loadPackagesFromImports + .mockReturnValueOnce( + new Promise((resolve) => { + finishFirstLoad = resolve; + }), + ) + .mockResolvedValueOnce(undefined); + const controller = new TestController(); + controller.setPyodide(pyodide); + + await startSession(controller, "first_dependency"); + await startSession(controller, "superseded_dependency"); + await controller.stopSession(); + await startSession(controller, "current_dependency"); + + expect(loadPackagesFromImports).toHaveBeenCalledOnce(); + finishFirstLoad(); + await vi.waitFor(() => + expect(loadPackagesFromImports).toHaveBeenCalledTimes(2), + ); + + const loadedSources = loadPackagesFromImports.mock.calls.map( + ([source]) => source as string, + ); + expect(loadedSources[1]).toContain("current_dependency"); + expect(loadedSources.join("\n")).not.toContain("superseded_dependency"); + }); +}); diff --git a/frontend/src/core/wasm/worker/bootstrap.ts b/frontend/src/core/wasm/worker/bootstrap.ts index 8339b422258..2b9318eab24 100644 --- a/frontend/src/core/wasm/worker/bootstrap.ts +++ b/frontend/src/core/wasm/worker/bootstrap.ts @@ -1,5 +1,6 @@ /* Copyright 2026 Marimo. All rights reserved. */ import { loadPyodide, type PyodideInterface } from "pyodide"; +import type { PyCallable, PyProxy } from "pyodide/ffi"; import type { UserConfig } from "@/core/config/config-schema"; import type { NotificationPayload } from "@/core/kernel/messages"; import type { JsonString } from "@/utils/json/base64"; @@ -12,6 +13,7 @@ import type { SerializedBridge, WasmController } from "./types"; import { shouldLoadDuckDBPackages } from "../utils"; const MAKE_SNAPSHOT = false; +type SessionResources = [PyProxy, PyCallable, PyProxy, PyCallable]; // This class initializes the wasm environment // We would like this initialization to be parallelizable @@ -25,6 +27,9 @@ const MAKE_SNAPSHOT = false; export class DefaultWasmController implements WasmController { protected pyodide: PyodideInterface | null = null; + private packageLoadQueue = Promise.resolve(); + private sessionGeneration = 0; + private stopCurrentSession: PyCallable | undefined; get requirePyodide() { invariant(this.pyodide, "Pyodide not loaded"); @@ -111,6 +116,7 @@ export class DefaultWasmController implements WasmController { onMessage: (message: JsonString) => void; userConfig: UserConfig; }): Promise { + const sessionGeneration = this.sessionGeneration; const { code, filename, onMessage, queryParameters, userConfig } = opts; // We pass down a messenger object to the code // This is used to have synchronous communication between the JS and Python code @@ -126,47 +132,120 @@ export class DefaultWasmController implements WasmController { const span = t.startSpan("startSession.runPython"); const nbFilename = filename || WasmFileSystem.NOTEBOOK_FILENAME; - const [bridge, init, packages] = this.requirePyodide.runPython( + const sessionResources = this.requirePyodide.runPython( ` print("[py] Starting marimo...") import asyncio + import gc import js from marimo._pyodide.bootstrap import create_session, instantiate assert js.messenger, "messenger is not defined" assert js.query_params, "query_params is not defined" - session, bridge = create_session( - filename="${nbFilename}", - query_params=js.query_params.to_py(), - message_callback=js.messenger.callback, - user_config=js.user_config.to_py(), - ) + def create_session_resources(): + session, bridge = create_session( + filename="${nbFilename}", + query_params=js.query_params.to_py(), + message_callback=js.messenger.callback, + user_config=js.user_config.to_py(), + ) + session_task = None - def init(auto_instantiate=True): - instantiate(session, auto_instantiate) - asyncio.create_task(session.start()) + def init(auto_instantiate=True): + nonlocal session_task + instantiate(session, auto_instantiate) + session_task = asyncio.create_task(session.start()) - # Find the packages to install - with open("${nbFilename}", "r") as f: - packages = session.find_packages(f.read()) + async def stop(): + nonlocal bridge, session, session_task + task = session_task + try: + kernel_task = getattr(session, "kernel_task", None) + if kernel_task is None: + if task is not None: + task.cancel() + else: + kernel_task.stop() + if task is not None: + try: + await task + except asyncio.CancelledError: + pass + finally: + session_task = None + session = None + bridge = None + gc.collect() - bridge, init, packages`, - ); - span.end(); + with open("${nbFilename}", "r") as f: + packages = session.find_packages(f.read()) - const foundPackages = new Set(packages.toJs()); + return bridge, init, packages, stop + + create_session_resources()`, + ) as PyProxy; + span.end(); + let bridgeProxy!: PyProxy; + let initSession!: PyCallable; + let packagesProxy!: PyProxy; + let stopSession!: PyCallable; + try { + [bridgeProxy, initSession, packagesProxy, stopSession] = + sessionResources as unknown as SessionResources; + } finally { + sessionResources.destroy(); + } + let foundPackages: Set; + try { + foundPackages = new Set(packagesProxy.toJs()); + } catch (error) { + bridgeProxy.destroy(); + initSession.destroy(); + stopSession.destroy(); + throw error; + } finally { + packagesProxy.destroy(); + } + this.stopCurrentSession?.destroy(); + this.stopCurrentSession = stopSession; // Fire and forget: // Load notebook dependencies and instantiate the session // We don't want to wait for this to finish, // so we can show the initial code immediately giving // a sense of responsiveness. - void this.loadNotebookDeps(code, foundPackages).then(() => { - return init(userConfig.runtime.auto_instantiate); + const dependenciesReady = this.packageLoadQueue.then(() => { + if (sessionGeneration !== this.sessionGeneration) { + return; + } + return this.loadNotebookDeps(code, foundPackages); }); + this.packageLoadQueue = dependenciesReady.catch(() => undefined); + void dependenciesReady + .then(() => { + if (sessionGeneration !== this.sessionGeneration) { + return; + } + return initSession(userConfig.runtime.auto_instantiate); + }) + .catch((error: unknown) => { + Logger.error("Failed to load notebook dependencies", error); + }) + .finally(() => initSession.destroy()); - return bridge; + return bridgeProxy as unknown as SerializedBridge; + } + + async stopSession(): Promise { + this.sessionGeneration += 1; + const stop = this.stopCurrentSession; + this.stopCurrentSession = undefined; + try { + await stop?.(); + } finally { + stop?.destroy(); + } } private async loadNotebookDeps(code: string, foundPackages: Set) { From 634d3a1eea8610d38c9d6221be9002eac80ccf31 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?P=C3=A9ter=20Ferenc=20Gyarmati?= Date: Sat, 11 Jul 2026 20:33:02 +0200 Subject: [PATCH 03/12] feat(islands): add retained app lifecycle Export canReplaceApp(), initialize(), and stopApp() for hosts that preserve the browser realm across page navigation. Reuse the global worker while serializing app replacement and scoping worker requests by app ID and session generation so stale work cannot reach the active app. --- .../src/core/islands/__tests__/bridge.test.ts | 139 +++++++++++- frontend/src/core/islands/bridge.ts | 113 ++++++++-- frontend/src/core/islands/main.ts | 30 ++- frontend/src/core/islands/worker/worker.tsx | 204 +++++++++++++----- 4 files changed, 413 insertions(+), 73 deletions(-) diff --git a/frontend/src/core/islands/__tests__/bridge.test.ts b/frontend/src/core/islands/__tests__/bridge.test.ts index cac5fa2e9ca..9a05e4076c2 100644 --- a/frontend/src/core/islands/__tests__/bridge.test.ts +++ b/frontend/src/core/islands/__tests__/bridge.test.ts @@ -45,6 +45,8 @@ vi.stubGlobal("URL", MockURL); const { mockBridge, mockLoadPackages, + mockReplaceSessionRequest, + mockStopSessionRequest, mockStartSessionRequest, mockParseMarimoIslandApps, mockCreateMarimoFile, @@ -52,6 +54,8 @@ const { } = vi.hoisted(() => ({ mockBridge: vi.fn(), mockLoadPackages: vi.fn(), + mockReplaceSessionRequest: vi.fn(), + mockStopSessionRequest: vi.fn(), mockStartSessionRequest: vi.fn(), mockParseMarimoIslandApps: vi.fn<() => TestIslandApp[]>(() => []), mockCreateMarimoFile: vi.fn(), @@ -66,7 +70,9 @@ vi.mock("@/core/wasm/rpc", () => ({ request: { bridge: mockBridge, loadPackages: mockLoadPackages, + replaceSession: mockReplaceSessionRequest, startSession: mockStartSessionRequest, + stopSession: mockStopSessionRequest, }, send: { consumerReady: vi.fn(), @@ -118,6 +124,17 @@ describe("IslandsPyodideBridge", () => { bridge = new IslandsPyodideBridge({ autoStartSessions: false }); }); + function mockSingleApp(file = "generated app 1") { + mockParseMarimoIslandApps.mockReturnValue([ + { + id: "app-1", + cells: [{ code: "x = 1", idx: 0, output: "
1
" }], + }, + ]); + mockCreateMarimoFile.mockReturnValue(file); + bridge.workerReady.resolve(); + } + describe("startSessionsForAllApps", () => { it("should prefer trusted export notebook code when there is exactly one reactive app", async () => { mockParseMarimoIslandApps.mockReturnValue([ @@ -140,6 +157,7 @@ describe("IslandsPyodideBridge", () => { expect(mockStartSessionRequest).toHaveBeenCalledWith({ appId: "app-1", code: "import marimo\napp = marimo.App()\n@app.cell\ndef __():\n x = 1\n return", + sessionGeneration: 1, }); }); @@ -164,6 +182,7 @@ describe("IslandsPyodideBridge", () => { expect(mockStartSessionRequest).toHaveBeenCalledWith({ appId: "app-1", code: "generated payload app", + sessionGeneration: 1, }); }); @@ -194,10 +213,32 @@ describe("IslandsPyodideBridge", () => { expect(mockStartSessionRequest).toHaveBeenNthCalledWith(1, { appId: "app-1", code: "generated app 1", + sessionGeneration: 1, }); expect(mockStartSessionRequest).toHaveBeenNthCalledWith(2, { appId: "app-2", code: "generated app 2", + sessionGeneration: 2, + }); + + await bridge.sendComponentValues({ + objectIds: [uiElementId("slider-1")], + values: [2], + }); + await bridge.stopSession(); + + expect(mockReplaceSessionRequest).not.toHaveBeenCalled(); + expect(mockStopSessionRequest).not.toHaveBeenCalled(); + expect(mockBridge).toHaveBeenCalledWith({ + appId: "app-1", + functionName: "put_control_request", + payload: { + objectIds: ["slider-1"], + token: "test-uuid-12345", + type: "update-ui-element", + values: [2], + }, + sessionGeneration: 1, }); }); @@ -218,10 +259,102 @@ describe("IslandsPyodideBridge", () => { expect(mockStartSessionRequest).toHaveBeenCalledWith({ appId: "app-1", code: "generated app 1", + sessionGeneration: 1, }); }); }); + describe("app lifecycle", () => { + it("replaces the active single app and skips duplicate initialization", async () => { + mockSingleApp(); + + await bridge.initializeApps(); + await bridge.initializeApps(); + await bridge.sendComponentValues({ + objectIds: [uiElementId("slider-1")], + values: [2], + }); + + expect(mockReplaceSessionRequest).toHaveBeenCalledOnce(); + expect(mockReplaceSessionRequest).toHaveBeenCalledWith({ + appId: "app-1", + code: "generated app 1", + sessionGeneration: 1, + }); + expect(mockBridge).toHaveBeenCalledWith( + expect.objectContaining({ + appId: "app-1", + sessionGeneration: 1, + }), + ); + mockCreateMarimoFile.mockReturnValue("generated app 2"); + await bridge.initializeApps(); + + expect(mockReplaceSessionRequest).toHaveBeenCalledTimes(2); + expect(mockReplaceSessionRequest).toHaveBeenNthCalledWith(2, { + appId: "app-1", + code: "generated app 2", + sessionGeneration: 2, + }); + }); + + it("stops the matching active app", async () => { + mockSingleApp(); + await bridge.initializeApps(); + + await bridge.stopSession("other-app"); + await bridge.stopSession("app-1"); + + expect(mockStopSessionRequest).toHaveBeenCalledOnce(); + expect(mockStopSessionRequest).toHaveBeenCalledWith({ + appId: "app-1", + sessionGeneration: 1, + }); + }); + + it("allows initialization to retry after replacement fails", async () => { + mockSingleApp(); + mockReplaceSessionRequest + .mockRejectedValueOnce(new Error("replacement failed")) + .mockResolvedValueOnce(undefined); + + await expect(bridge.initializeApps()).rejects.toThrow( + "replacement failed", + ); + await bridge.initializeApps(); + + expect(mockReplaceSessionRequest).toHaveBeenCalledTimes(2); + }); + + it("scopes controls while the replacement session is pending", async () => { + let finishReplacement!: () => void; + mockReplaceSessionRequest.mockReturnValueOnce( + new Promise((resolve) => { + finishReplacement = resolve; + }), + ); + mockSingleApp(); + + const initialization = bridge.initializeApps(); + await vi.waitFor(() => + expect(mockReplaceSessionRequest).toHaveBeenCalledOnce(), + ); + await bridge.sendComponentValues({ + objectIds: [uiElementId("slider-1")], + values: [2], + }); + + expect(mockBridge).toHaveBeenCalledWith( + expect.objectContaining({ + appId: "app-1", + sessionGeneration: 1, + }), + ); + finishReplacement(); + await initialization; + }); + }); + describe("sendComponentValues", () => { it("should include type field and token in control request", async () => { const request = { @@ -313,7 +446,11 @@ describe("IslandsPyodideBridge", () => { await bridge.sendRun(request); // Verify loadPackages was called with joined codes - expect(mockLoadPackages).toHaveBeenCalledWith("import pandas"); + expect(mockLoadPackages).toHaveBeenCalledWith({ + appId: undefined, + code: "import pandas", + sessionGeneration: undefined, + }); // Verify order: loadPackages should be called before bridge const loadPackagesCallOrder = diff --git a/frontend/src/core/islands/bridge.ts b/frontend/src/core/islands/bridge.ts index e3ed684d3f0..842b1ce336b 100644 --- a/frontend/src/core/islands/bridge.ts +++ b/frontend/src/core/islands/bridge.ts @@ -63,8 +63,14 @@ export class IslandsPyodideBridge implements RunRequests, EditRequests { private readonly store: typeof defaultStore; private readonly root: Document | Element; private readonly autoStartSessions: boolean; + private session: + | { appId: string; code?: string; sessionGeneration: number } + | undefined; + private nextSessionGeneration = 0; + private appTransition = Promise.resolve(); public initialized = new Deferred(); + public workerReady = new Deferred(); constructor(config: IslandsBridgeConfig = {}) { this.store = config.store || defaultStore; @@ -88,8 +94,9 @@ export class IslandsPyodideBridge implements RunRequests, EditRequests { */ private setupMessageListeners(): void { this.rpc.addMessageListener("ready", () => { + this.workerReady.resolve(); if (this.autoStartSessions) { - this.startSessionsForAllApps(); + void this.startSessionsForAllApps(); } }); @@ -118,8 +125,18 @@ export class IslandsPyodideBridge implements RunRequests, EditRequests { /** * Starts sessions for all apps found in the DOM */ - private startSessionsForAllApps(): void { + private startSessionsForAllApps(): Promise { + return this.enqueueAppTransition(() => this.startApps(false)); + } + + async initializeApps(): Promise { + await this.workerReady.promise; + await this.enqueueAppTransition(() => this.startApps(true)); + } + + private async startApps(replaceSingleApp: boolean): Promise { const apps = parseMarimoIslandApps(this.root); + const replacesSession = replaceSingleApp && apps.length === 1; Logger.debug( `Starting sessions for ${apps.length} app(s):`, apps.map((a) => `${a.id} (${a.cells.length} cells)`), @@ -134,20 +151,75 @@ export class IslandsPyodideBridge implements RunRequests, EditRequests { const notebookCode = exportContext?.notebookCode; for (const app of apps) { const file = notebookCode || createMarimoFile(app); + if ( + replacesSession && + this.session?.appId === app.id && + this.session.code === file + ) { + return; + } Logger.debug(`App ${app.id} marimo file:\n`, file); - this.startSession({ + const request = { code: file, appId: app.id, - }).catch((error) => { + sessionGeneration: ++this.nextSessionGeneration, + }; + const previousSession = this.session; + if (replacesSession || !previousSession) { + this.session = { + ...request, + code: replacesSession ? file : undefined, + }; + } + const operation = replacesSession + ? this.rpc.proxy.request.replaceSession(request) + : this.startSession(request); + try { + await operation; + } catch (error) { + if (this.session?.sessionGeneration === request.sessionGeneration) { + this.session = replacesSession ? undefined : previousSession; + } Logger.error(`Failed to start session for app ${app.id}:`, error); - }); + if (replacesSession) { + throw error; + } + } } } + async stopSession(appId?: string): Promise { + await this.enqueueAppTransition(async () => { + const session = this.session; + if (session?.code === undefined || (appId && session.appId !== appId)) { + return; + } + try { + await this.rpc.proxy.request.stopSession( + appId === undefined + ? { sessionGeneration: session.sessionGeneration } + : { appId, sessionGeneration: session.sessionGeneration }, + ); + } finally { + this.session = undefined; + } + }); + } + + private enqueueAppTransition(operation: () => Promise): Promise { + const result = this.appTransition.then(operation); + this.appTransition = result.catch(() => undefined); + return result; + } + /** * Starts a new Python session for an app */ - async startSession(opts: { code: string; appId: string }): Promise { + async startSession(opts: { + code: string; + appId: string; + sessionGeneration: number; + }): Promise { await this.rpc.proxy.request.startSession(opts); } @@ -203,11 +275,19 @@ export class IslandsPyodideBridge implements RunRequests, EditRequests { // ============================================================================ sendRun: EditRequests["sendRun"] = async (request): Promise => { - await this.rpc.proxy.request.loadPackages(request.codes.join("\n")); - await this.putControlRequest({ - type: "execute-cells", - ...request, + const session = this.session; + await this.rpc.proxy.request.loadPackages({ + appId: session?.appId, + code: request.codes.join("\n"), + sessionGeneration: session?.sessionGeneration, }); + await this.putControlRequest( + { + type: "execute-cells", + ...request, + }, + session, + ); return null; }; @@ -278,10 +358,17 @@ export class IslandsPyodideBridge implements RunRequests, EditRequests { // The kernel uses msgspec to parse control requests, which requires a 'type' // field for discriminated union deserialization. - private async putControlRequest(operation: CommandMessage): Promise { + private async putControlRequest( + operation: CommandMessage, + session = this.session, + ): Promise { await this.rpc.proxy.request.bridge({ functionName: "put_control_request", payload: operation, + ...(session && { + appId: session.appId, + sessionGeneration: session.sessionGeneration, + }), }); } @@ -301,7 +388,9 @@ let globalBridgeInstance: IslandsPyodideBridge | null = null; export function getGlobalBridge(): IslandsPyodideBridge { if (!globalBridgeInstance) { - globalBridgeInstance = new IslandsPyodideBridge(); + globalBridgeInstance = new IslandsPyodideBridge({ + autoStartSessions: false, + }); } return globalBridgeInstance; } diff --git a/frontend/src/core/islands/main.ts b/frontend/src/core/islands/main.ts index 21dd168b8db..d441cba0c18 100644 --- a/frontend/src/core/islands/main.ts +++ b/frontend/src/core/islands/main.ts @@ -14,12 +14,36 @@ import "iconify-icon"; import { Logger } from "@/utils/Logger"; import { initializeIslands } from "./bootstrap"; import { getGlobalBridge } from "./bridge"; +import { ISLAND_TAG_NAMES } from "./constants"; +import { parseMarimoIslandApps } from "./parse"; + +const bridge = getGlobalBridge(); +let bootstrapPromise: Promise | undefined; + +/** Returns whether the current document can replace its app in this worker. */ +export function canReplaceApp(): boolean { + return parseMarimoIslandApps(document, { materialize: false }).length === 1; +} /** - * Main entry point for the js bundle for embedded marimo apps. + * Mounts marimo custom elements and starts the apps in the current document. */ -export async function initialize() { - await initializeIslands({ bridge: getGlobalBridge() }); +export async function initialize(): Promise { + if (!document.querySelector(ISLAND_TAG_NAMES.ISLAND)) { + return; + } + bootstrapPromise ??= initializeIslands({ bridge }).catch((error: unknown) => { + bootstrapPromise = undefined; + throw error; + }); + await bootstrapPromise; + await new Promise((resolve) => setTimeout(resolve, 0)); + await bridge.initializeApps(); +} + +/** Stops the matching active app session. */ +export async function stopApp(appId?: string): Promise { + await bridge.stopSession(appId); } // Auto-initialize on module load diff --git a/frontend/src/core/islands/worker/worker.tsx b/frontend/src/core/islands/worker/worker.tsx index 090236f9b86..83bd71d8837 100644 --- a/frontend/src/core/islands/worker/worker.tsx +++ b/frontend/src/core/islands/worker/worker.tsx @@ -1,6 +1,7 @@ /* Copyright 2026 Marimo. All rights reserved. */ import type { PyodideInterface } from "pyodide"; +import type { PyProxy } from "pyodide/ffi"; import { createRPC, createRPCRequestHandler, @@ -16,7 +17,6 @@ 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 { Deferred } from "../../../utils/Deferred"; import { prettyError } from "../../../utils/errors"; import { invariant } from "../../../utils/invariant"; import { ReadonlyWasmController } from "./controller"; @@ -50,59 +50,129 @@ const messageBuffer = new MessageBuffer( rpc.send.kernelMessage({ message }); }, ); -const bridgeReady = new Deferred(); +interface SessionRequest { + code: string; + appId: string; + sessionGeneration: number; +} + +let activeSession: + | (Omit & { bridge: SerializedBridge }) + | undefined; +let sessionQueue = Promise.resolve(); + +async function startSession( + opts: SessionRequest, + replace: boolean, +): Promise { + await pyodideReadyPromise; + + try { + invariant(self.controller, "Controller not loaded"); + if (replace) { + await stopActiveSession(); + } + const notebook = await self.controller.mountFilesystem({ + code: opts.code, + filename: `app-${opts.appId}.py`, + }); + const nextBridge = await self.controller.startSession({ + ...notebook, + onMessage: messageBuffer.push, + }); + if (activeSession && !replace) { + (nextBridge as unknown as PyProxy).destroy(); + } else { + (activeSession?.bridge as unknown as PyProxy | undefined)?.destroy(); + activeSession = { + appId: opts.appId, + bridge: nextBridge, + sessionGeneration: opts.sessionGeneration, + }; + } + rpc.send.initialized({}); + } catch (error) { + rpc.send.initializedError({ + error: prettyError(error), + }); + throw error; + } +} + +async function stopActiveSession(): Promise { + const session = activeSession; + activeSession = undefined; + (session?.bridge as unknown as PyProxy | undefined)?.destroy(); + await self.controller.stopSession(); +} + +function enqueueSession(operation: () => Promise): Promise { + const result = sessionQueue.then(operation); + sessionQueue = result.then( + () => undefined, + () => undefined, + ); + return result; +} // Handle RPC requests const requestHandler = createRPCRequestHandler({ /** * Start the session */ - startSession: async (opts: { code: string; appId: string }) => { - await pyodideReadyPromise; // Make sure loading is done + startSession: async (opts: SessionRequest) => { + await enqueueSession(() => startSession(opts, false)); + }, - try { + replaceSession: async (opts: SessionRequest) => { + await enqueueSession(() => startSession(opts, true)); + }, + + stopSession: async (opts: { appId?: string; sessionGeneration: number }) => { + await enqueueSession(async () => { + if ( + (opts.appId && activeSession?.appId !== opts.appId) || + activeSession?.sessionGeneration !== opts.sessionGeneration + ) { + return; + } invariant(self.controller, "Controller not loaded"); - const notebook = await self.controller.mountFilesystem({ - code: opts.code, - filename: `app-${opts.appId}.py`, - }); - const bridge = await self.controller.startSession({ - ...notebook, - onMessage: messageBuffer.push, - }); - bridgeReady.resolve(bridge); - rpc.send.initialized({}); - } catch (error) { - rpc.send.initializedError({ - error: prettyError(error), - }); - } - return; + await stopActiveSession(); + }); }, /** * Load packages */ - loadPackages: async (code: string) => { + loadPackages: async (opts: { + appId?: string; + code: string; + sessionGeneration?: number; + }) => { await pyodideReadyPromise; // Make sure loading is done + await enqueueSession(async () => { + requireActiveBridge(opts); + + let { code } = opts; - if (shouldLoadDuckDBPackages(code)) { - // Add pandas and duckdb to the code for mo.sql and for remote duckdb sources - code = `import pandas\n${code}`; - code = `import duckdb\n${code}`; - code = `import sqlglot\n${code}`; - - // Polars + SQL requires pyarrow, and installing - // after notebook load does not work. As a heuristic, - // if it appears that the notebook uses polars, add pyarrow. - if (code.includes("polars")) { - code = `import pyarrow\n${code}`; + if (shouldLoadDuckDBPackages(code)) { + // Add pandas and duckdb to the code for mo.sql and for remote duckdb sources + code = `import pandas\n${code}`; + code = `import duckdb\n${code}`; + code = `import sqlglot\n${code}`; + + // Polars + SQL requires pyarrow, and installing + // after notebook load does not work. As a heuristic, + // if it appears that the notebook uses polars, add pyarrow. + if (code.includes("polars")) { + code = `import pyarrow\n${code}`; + } } - } - await self.pyodide.loadPackagesFromImports(code, { - messageCallback: Logger.log, - errorCallback: Logger.error, + await self.pyodide.loadPackagesFromImports(code, { + messageCallback: Logger.log, + errorCallback: Logger.error, + }); }); }, @@ -110,37 +180,57 @@ const requestHandler = createRPCRequestHandler({ * Call a function on the bridge */ bridge: async (opts: { + appId?: string; functionName: keyof RawBridge; payload: {} | undefined | null; + sessionGeneration?: number; }) => { await pyodideReadyPromise; // Make sure loading is done const { functionName, payload } = opts; // Perform the function call to the Python bridge - const bridge = await bridgeReady.promise; - - // Serialize the payload - const payloadString = - payload == null - ? null - : typeof payload === "string" - ? payload - : JSON.stringify(payload); - - // Make the request - const response = - payloadString == null - ? // @ts-expect-error ehh TypeScript - await bridge[functionName]() - : // @ts-expect-error ehh TypeScript - await bridge[functionName](payloadString); - - // Post the response back to the main thread - return typeof response === "string" ? JSON.parse(response) : response; + return enqueueSession(async () => { + const bridge = requireActiveBridge(opts); + + // Serialize the payload + const payloadString = + payload == null + ? null + : typeof payload === "string" + ? payload + : JSON.stringify(payload); + + // Make the request + const response = + payloadString == null + ? // @ts-expect-error ehh TypeScript + await bridge[functionName]() + : // @ts-expect-error ehh TypeScript + await bridge[functionName](payloadString); + + // Post the response back to the main thread + return typeof response === "string" ? JSON.parse(response) : response; + }); }, }); +function requireActiveBridge(opts: { + appId?: string; + sessionGeneration?: number; +}): SerializedBridge { + const session = activeSession; + if ( + !session || + !opts.appId || + opts.appId !== session.appId || + opts.sessionGeneration !== session.sessionGeneration + ) { + throw new Error("The marimo app session is no longer active."); + } + return session.bridge; +} + // create the iframe's schema export type WorkerSchema = RPCSchema< { From 913aac1e2dab5401a79566d3fd25c1dd2ddccefa Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?P=C3=A9ter=20Ferenc=20Gyarmati?= Date: Wed, 15 Jul 2026 08:27:24 +0200 Subject: [PATCH 04/12] feat(islands): make lifecycle initialization repeatable --- frontend/src/core/islands/bootstrap.ts | 4 +++- frontend/src/core/islands/main.ts | 4 +++- 2 files changed, 6 insertions(+), 2 deletions(-) diff --git a/frontend/src/core/islands/bootstrap.ts b/frontend/src/core/islands/bootstrap.ts index be1eaa0086e..c60dbee0d40 100644 --- a/frontend/src/core/islands/bootstrap.ts +++ b/frontend/src/core/islands/bootstrap.ts @@ -208,7 +208,9 @@ function handleMessage( setKernelState: Functions.NOOP, onError: Logger.error, }); - defineCustomElement(ISLAND_TAG_NAMES.ISLAND, MarimoIslandElement); + if (!window.customElements?.get(ISLAND_TAG_NAMES.ISLAND)) { + defineCustomElement(ISLAND_TAG_NAMES.ISLAND, MarimoIslandElement); + } return; case "send-ui-element-message": diff --git a/frontend/src/core/islands/main.ts b/frontend/src/core/islands/main.ts index d441cba0c18..3dd9c57c828 100644 --- a/frontend/src/core/islands/main.ts +++ b/frontend/src/core/islands/main.ts @@ -22,6 +22,9 @@ let bootstrapPromise: Promise | undefined; /** Returns whether the current document can replace its app in this worker. */ export function canReplaceApp(): boolean { + if (!document.querySelector(ISLAND_TAG_NAMES.ISLAND)) { + return false; + } return parseMarimoIslandApps(document, { materialize: false }).length === 1; } @@ -37,7 +40,6 @@ export async function initialize(): Promise { throw error; }); await bootstrapPromise; - await new Promise((resolve) => setTimeout(resolve, 0)); await bridge.initializeApps(); } From abd9f55c7685199823eb74acd5c3e306f2526664 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?P=C3=A9ter=20Ferenc=20Gyarmati?= Date: Wed, 15 Jul 2026 08:27:30 +0200 Subject: [PATCH 05/12] feat(islands): scope controls to app transitions --- .../src/core/islands/__tests__/bridge.test.ts | 206 ++++++++++++++++-- frontend/src/core/islands/bridge.ts | 49 +++-- 2 files changed, 221 insertions(+), 34 deletions(-) diff --git a/frontend/src/core/islands/__tests__/bridge.test.ts b/frontend/src/core/islands/__tests__/bridge.test.ts index 9a05e4076c2..9fe9d83efd1 100644 --- a/frontend/src/core/islands/__tests__/bridge.test.ts +++ b/frontend/src/core/islands/__tests__/bridge.test.ts @@ -135,6 +135,11 @@ describe("IslandsPyodideBridge", () => { bridge.workerReady.resolve(); } + async function initializeSingleApp() { + mockSingleApp(); + await bridge.initializeApps(); + } + describe("startSessionsForAllApps", () => { it("should prefer trusted export notebook code when there is exactly one reactive app", async () => { mockParseMarimoIslandApps.mockReturnValue([ @@ -265,7 +270,35 @@ describe("IslandsPyodideBridge", () => { }); describe("app lifecycle", () => { - it("replaces the active single app and skips duplicate initialization", async () => { + it("holds controls until the first app session is ready", async () => { + mockParseMarimoIslandApps.mockReturnValue([ + { + id: "app-1", + cells: [{ code: "x = 1", idx: 0, output: "
1
" }], + }, + ]); + mockCreateMarimoFile.mockReturnValue("generated app 1"); + + const initialization = bridge.initializeApps(); + const control = bridge.sendComponentValues({ + objectIds: [uiElementId("slider-1")], + values: [2], + }); + await Promise.resolve(); + + expect(mockBridge).not.toHaveBeenCalled(); + + bridge.workerReady.resolve(); + await Promise.all([initialization, control]); + expect(mockBridge).toHaveBeenCalledWith( + expect.objectContaining({ + appId: "app-1", + sessionGeneration: 1, + }), + ); + }); + + it("starts the first app, skips duplicates, and replaces changed source", async () => { mockSingleApp(); await bridge.initializeApps(); @@ -275,12 +308,13 @@ describe("IslandsPyodideBridge", () => { values: [2], }); - expect(mockReplaceSessionRequest).toHaveBeenCalledOnce(); - expect(mockReplaceSessionRequest).toHaveBeenCalledWith({ + expect(mockStartSessionRequest).toHaveBeenCalledOnce(); + expect(mockStartSessionRequest).toHaveBeenCalledWith({ appId: "app-1", code: "generated app 1", sessionGeneration: 1, }); + expect(mockReplaceSessionRequest).not.toHaveBeenCalled(); expect(mockBridge).toHaveBeenCalledWith( expect.objectContaining({ appId: "app-1", @@ -290,8 +324,8 @@ describe("IslandsPyodideBridge", () => { mockCreateMarimoFile.mockReturnValue("generated app 2"); await bridge.initializeApps(); - expect(mockReplaceSessionRequest).toHaveBeenCalledTimes(2); - expect(mockReplaceSessionRequest).toHaveBeenNthCalledWith(2, { + expect(mockReplaceSessionRequest).toHaveBeenCalledOnce(); + expect(mockReplaceSessionRequest).toHaveBeenCalledWith({ appId: "app-1", code: "generated app 2", sessionGeneration: 2, @@ -312,50 +346,162 @@ describe("IslandsPyodideBridge", () => { }); }); + it("keeps controls sent during stop scoped to the stopped app", async () => { + mockSingleApp(); + await bridge.initializeApps(); + mockBridge.mockClear(); + let finishStop!: () => void; + mockStopSessionRequest.mockReturnValueOnce( + new Promise((resolve) => { + finishStop = resolve; + }), + ); + + const stop = bridge.stopSession("app-1"); + await vi.waitFor(() => + expect(mockStopSessionRequest).toHaveBeenCalledOnce(), + ); + const control = bridge.sendComponentValues({ + objectIds: [uiElementId("slider-1")], + values: [2], + }); + mockParseMarimoIslandApps.mockReturnValue([ + { + id: "app-2", + cells: [{ code: "x = 2", idx: 0, output: "
2
" }], + }, + ]); + mockCreateMarimoFile.mockReturnValue("generated app 2"); + const nextInitialization = bridge.initializeApps(); + + finishStop(); + await Promise.all([stop, control, nextInitialization]); + + expect(mockBridge).toHaveBeenCalledWith( + expect.objectContaining({ + appId: "app-1", + sessionGeneration: 1, + }), + ); + }); + it("allows initialization to retry after replacement fails", async () => { mockSingleApp(); - mockReplaceSessionRequest - .mockRejectedValueOnce(new Error("replacement failed")) - .mockResolvedValueOnce(undefined); + await bridge.initializeApps(); + mockCreateMarimoFile.mockReturnValue("generated app 2"); + mockReplaceSessionRequest.mockRejectedValueOnce( + new Error("replacement failed"), + ); await expect(bridge.initializeApps()).rejects.toThrow( "replacement failed", ); await bridge.initializeApps(); - expect(mockReplaceSessionRequest).toHaveBeenCalledTimes(2); + expect(mockReplaceSessionRequest).toHaveBeenCalledOnce(); + expect(mockStartSessionRequest).toHaveBeenCalledTimes(2); + expect(mockStartSessionRequest).toHaveBeenLastCalledWith({ + appId: "app-1", + code: "generated app 2", + sessionGeneration: 3, + }); }); - it("scopes controls while the replacement session is pending", async () => { - let finishReplacement!: () => void; - mockReplaceSessionRequest.mockReturnValueOnce( - new Promise((resolve) => { - finishReplacement = resolve; + it("scopes controls to the replacement active when they were sent", async () => { + mockSingleApp(); + await bridge.initializeApps(); + mockBridge.mockClear(); + mockCreateMarimoFile.mockReturnValue("generated app 2"); + let finishSecondApp!: () => void; + let finishThirdApp!: () => void; + mockReplaceSessionRequest + .mockReturnValueOnce( + new Promise((resolve) => { + finishSecondApp = resolve; + }), + ) + .mockReturnValueOnce( + new Promise((resolve) => { + finishThirdApp = resolve; + }), + ); + + const secondInitialization = bridge.initializeApps(); + await vi.waitFor(() => + expect(mockReplaceSessionRequest).toHaveBeenCalledOnce(), + ); + const control = bridge.sendComponentValues({ + objectIds: [uiElementId("slider-1")], + values: [2], + }); + mockParseMarimoIslandApps.mockReturnValue([ + { + id: "app-2", + cells: [{ code: "x = 3", idx: 0, output: "
3
" }], + }, + ]); + mockCreateMarimoFile.mockReturnValue("generated app 3"); + const thirdInitialization = bridge.initializeApps(); + + expect(mockBridge).not.toHaveBeenCalled(); + + finishSecondApp(); + await vi.waitFor(() => expect(mockBridge).toHaveBeenCalledOnce()); + expect(mockBridge).toHaveBeenCalledWith( + expect.objectContaining({ + appId: "app-1", + sessionGeneration: 2, }), ); + + finishThirdApp(); + await Promise.all([secondInitialization, thirdInitialization, control]); + }); + + it("keeps controls from a failed replacement out of the next app", async () => { mockSingleApp(); + await bridge.initializeApps(); + mockBridge.mockClear(); + mockCreateMarimoFile.mockReturnValue("generated app 2"); + let failReplacement!: (error: Error) => void; + mockReplaceSessionRequest.mockReturnValueOnce( + new Promise((_resolve, reject) => { + failReplacement = reject; + }), + ); - const initialization = bridge.initializeApps(); + const failedInitialization = bridge.initializeApps(); await vi.waitFor(() => expect(mockReplaceSessionRequest).toHaveBeenCalledOnce(), ); - await bridge.sendComponentValues({ + const control = bridge.sendComponentValues({ objectIds: [uiElementId("slider-1")], values: [2], }); + failReplacement(new Error("replacement failed")); + await expect(failedInitialization).rejects.toThrow("replacement failed"); + + mockParseMarimoIslandApps.mockReturnValue([ + { + id: "app-2", + cells: [{ code: "x = 3", idx: 0, output: "
3
" }], + }, + ]); + mockCreateMarimoFile.mockReturnValue("generated app 3"); + await Promise.all([bridge.initializeApps(), control]); expect(mockBridge).toHaveBeenCalledWith( expect.objectContaining({ appId: "app-1", - sessionGeneration: 1, + sessionGeneration: 2, }), ); - finishReplacement(); - await initialization; }); }); describe("sendComponentValues", () => { + beforeEach(initializeSingleApp); + it("should include type field and token in control request", async () => { const request = { objectIds: [uiElementId("Hbol-0")], @@ -365,6 +511,7 @@ describe("IslandsPyodideBridge", () => { await bridge.sendComponentValues(request); expect(mockBridge).toHaveBeenCalledWith({ + appId: "app-1", functionName: "put_control_request", payload: { type: "update-ui-element", @@ -372,6 +519,7 @@ describe("IslandsPyodideBridge", () => { values: [58], token: "test-uuid-12345", }, + sessionGeneration: 1, }); }); @@ -384,17 +532,21 @@ describe("IslandsPyodideBridge", () => { await bridge.sendComponentValues(request); expect(mockBridge).toHaveBeenCalledWith({ + appId: "app-1", functionName: "put_control_request", payload: expect.objectContaining({ type: "update-ui-element", objectIds: ["slider-1", "slider-2"], values: [10, 20], }), + sessionGeneration: 1, }); }); }); describe("sendFunctionRequest", () => { + beforeEach(initializeSingleApp); + it("should include type field in control request", async () => { const request = { functionCallId: requestId("call-123"), @@ -406,6 +558,7 @@ describe("IslandsPyodideBridge", () => { await bridge.sendFunctionRequest(request); expect(mockBridge).toHaveBeenCalledWith({ + appId: "app-1", functionName: "put_control_request", payload: { type: "invoke-function", @@ -414,11 +567,14 @@ describe("IslandsPyodideBridge", () => { functionName: "my_function", args: { x: 1, y: 2 }, }, + sessionGeneration: 1, }); }); }); describe("sendRun", () => { + beforeEach(initializeSingleApp); + it("should include type field in control request", async () => { const request = { cellIds: [cellId("cell-1"), cellId("cell-2")], @@ -428,12 +584,14 @@ describe("IslandsPyodideBridge", () => { await bridge.sendRun(request); expect(mockBridge).toHaveBeenCalledWith({ + appId: "app-1", functionName: "put_control_request", payload: { type: "execute-cells", cellIds: ["cell-1", "cell-2"], codes: ["print('hello')", "print('world')"], }, + sessionGeneration: 1, }); }); @@ -447,9 +605,9 @@ describe("IslandsPyodideBridge", () => { // Verify loadPackages was called with joined codes expect(mockLoadPackages).toHaveBeenCalledWith({ - appId: undefined, + appId: "app-1", code: "import pandas", - sessionGeneration: undefined, + sessionGeneration: 1, }); // Verify order: loadPackages should be called before bridge @@ -461,6 +619,8 @@ describe("IslandsPyodideBridge", () => { }); describe("sendModelValue", () => { + beforeEach(initializeSingleApp); + it("should include type field in control request", async () => { const request = { modelId: widgetModelId("widget-1"), @@ -475,6 +635,7 @@ describe("IslandsPyodideBridge", () => { await bridge.sendModelValue(request); expect(mockBridge).toHaveBeenCalledWith({ + appId: "app-1", functionName: "put_control_request", payload: { type: "model", @@ -486,11 +647,14 @@ describe("IslandsPyodideBridge", () => { }, buffers: [], }, + sessionGeneration: 1, }); }); }); describe("control request message format", () => { + beforeEach(initializeSingleApp); + it("should always include the type field required by msgspec", async () => { // Test all methods to ensure they include the type field await bridge.sendComponentValues({ objectIds: [], values: [] }); diff --git a/frontend/src/core/islands/bridge.ts b/frontend/src/core/islands/bridge.ts index 842b1ce336b..f077823809c 100644 --- a/frontend/src/core/islands/bridge.ts +++ b/frontend/src/core/islands/bridge.ts @@ -42,6 +42,12 @@ export interface IslandsBridgeConfig { autoStartSessions?: boolean; } +interface AppSession { + appId: string; + code?: string; + sessionGeneration: number; +} + /** * Bridge between the browser and Pyodide worker for islands mode. * @@ -63,9 +69,8 @@ export class IslandsPyodideBridge implements RunRequests, EditRequests { private readonly store: typeof defaultStore; private readonly root: Document | Element; private readonly autoStartSessions: boolean; - private session: - | { appId: string; code?: string; sessionGeneration: number } - | undefined; + private session: AppSession | undefined; + private sessionReady = new Deferred(); private nextSessionGeneration = 0; private appTransition = Promise.resolve(); @@ -130,13 +135,15 @@ export class IslandsPyodideBridge implements RunRequests, EditRequests { } async initializeApps(): Promise { - await this.workerReady.promise; - await this.enqueueAppTransition(() => this.startApps(true)); + await this.enqueueAppTransition(async () => { + await this.workerReady.promise; + await this.startApps(true); + }); } private async startApps(replaceSingleApp: boolean): Promise { const apps = parseMarimoIslandApps(this.root); - const replacesSession = replaceSingleApp && apps.length === 1; + const managesSingleApp = replaceSingleApp && apps.length === 1; Logger.debug( `Starting sessions for ${apps.length} app(s):`, apps.map((a) => `${a.id} (${a.cells.length} cells)`), @@ -152,7 +159,7 @@ export class IslandsPyodideBridge implements RunRequests, EditRequests { for (const app of apps) { const file = notebookCode || createMarimoFile(app); if ( - replacesSession && + managesSingleApp && this.session?.appId === app.id && this.session.code === file ) { @@ -165,10 +172,14 @@ export class IslandsPyodideBridge implements RunRequests, EditRequests { sessionGeneration: ++this.nextSessionGeneration, }; const previousSession = this.session; - if (replacesSession || !previousSession) { + const replacesSession = managesSingleApp && previousSession !== undefined; + if (managesSingleApp || !previousSession) { + if (managesSingleApp && this.sessionReady.status !== "pending") { + this.sessionReady = new Deferred(); + } this.session = { ...request, - code: replacesSession ? file : undefined, + code: managesSingleApp ? file : undefined, }; } const operation = replacesSession @@ -176,12 +187,15 @@ export class IslandsPyodideBridge implements RunRequests, EditRequests { : this.startSession(request); try { await operation; + if (this.sessionReady.status === "pending" && this.session) { + this.sessionReady.resolve(this.session); + } } catch (error) { if (this.session?.sessionGeneration === request.sessionGeneration) { - this.session = replacesSession ? undefined : previousSession; + this.session = managesSingleApp ? undefined : previousSession; } Logger.error(`Failed to start session for app ${app.id}:`, error); - if (replacesSession) { + if (managesSingleApp) { throw error; } } @@ -275,7 +289,7 @@ export class IslandsPyodideBridge implements RunRequests, EditRequests { // ============================================================================ sendRun: EditRequests["sendRun"] = async (request): Promise => { - const session = this.session; + const session = await this.getActiveSession(); await this.rpc.proxy.request.loadPackages({ appId: session?.appId, code: request.codes.join("\n"), @@ -360,8 +374,9 @@ export class IslandsPyodideBridge implements RunRequests, EditRequests { // field for discriminated union deserialization. private async putControlRequest( operation: CommandMessage, - session = this.session, + session?: AppSession, ): Promise { + session ??= await this.getActiveSession(); await this.rpc.proxy.request.bridge({ functionName: "put_control_request", payload: operation, @@ -372,6 +387,14 @@ export class IslandsPyodideBridge implements RunRequests, EditRequests { }); } + private async getActiveSession(): Promise { + const transition = this.appTransition; + const session = this.session; + const ready = this.sessionReady; + await transition; + return this.session ?? session ?? ready.promise; + } + /** * Cleans up resources (for testing) */ From 70e0966f2e14c7f8ac9af3bba7270a9e091a05d0 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?P=C3=A9ter=20Ferenc=20Gyarmati?= Date: Wed, 15 Jul 2026 08:27:35 +0200 Subject: [PATCH 06/12] feat(islands): refresh reused island source --- .../src/core/islands/__tests__/parse.test.ts | 17 +++++++++++++++++ frontend/src/core/islands/parse.ts | 18 +++++++----------- 2 files changed, 24 insertions(+), 11 deletions(-) diff --git a/frontend/src/core/islands/__tests__/parse.test.ts b/frontend/src/core/islands/__tests__/parse.test.ts index 547f429b2f5..18f6f3197d0 100644 --- a/frontend/src/core/islands/__tests__/parse.test.ts +++ b/frontend/src/core/islands/__tests__/parse.test.ts @@ -331,6 +331,23 @@ describe("parseIslandElement", () => { output: "
retained output
", }); }); + + it("prefers new live source when an island element is reused", () => { + const element = createMockIslandElement({ + code: 'print("new")', + innerHTML: "
new output
", + }); + element.setAttribute(ISLAND_DATA_ATTRIBUTES.REACTIVE, "true"); + retainIslandSource(element, { + code: 'print("retained")', + output: "
retained output
", + }); + + expect(parseIslandElement(element)).toEqual({ + code: 'print("new")', + output: "
new output
", + }); + }); }); describe("parseIslandElementsIntoApps", () => { diff --git a/frontend/src/core/islands/parse.ts b/frontend/src/core/islands/parse.ts index a0f83725d1c..b163627cd0e 100644 --- a/frontend/src/core/islands/parse.ts +++ b/frontend/src/core/islands/parse.ts @@ -340,24 +340,20 @@ export function retainIslandSource( export function parseIslandElement( embed: HTMLElement, ): { output: string; code: string } | null { - const retained = retainedIslandSources.get(embed); - if (retained) { - return retained.code ? retained : null; - } - const cellOutput = embed.querySelector( ISLAND_TAG_NAMES.CELL_OUTPUT, ); 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: { From 77c1b1cb2fc9f4c12cd4bc0c83d300d00be5002f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?P=C3=A9ter=20Ferenc=20Gyarmati?= Date: Fri, 17 Jul 2026 12:59:58 +0200 Subject: [PATCH 07/12] feat(wasm): make retained session teardown retryable --- .../worker/__tests__/controller.test.ts | 75 +++++++++++++++++-- frontend/src/core/wasm/worker/bootstrap.ts | 55 +++++++------- 2 files changed, 97 insertions(+), 33 deletions(-) diff --git a/frontend/src/core/islands/worker/__tests__/controller.test.ts b/frontend/src/core/islands/worker/__tests__/controller.test.ts index a583a292443..b86a8a646e0 100644 --- a/frontend/src/core/islands/worker/__tests__/controller.test.ts +++ b/frontend/src/core/islands/worker/__tests__/controller.test.ts @@ -10,30 +10,43 @@ class TestController extends ReadonlyWasmController { } } -function createPyodideStub() { +function createSessionResources( + stopImplementation: () => Promise = () => Promise.resolve(), +) { + const bridge = { destroy: vi.fn() }; const init = Object.assign(vi.fn(), { destroy: vi.fn() }); - const stop = Object.assign(vi.fn().mockResolvedValue(undefined), { + const stop = Object.assign(vi.fn(stopImplementation), { destroy: vi.fn(), }); const packages = { destroy: vi.fn(), toJs: () => [] }; - const sessionResources = Object.assign([{}, init, packages, stop], { + const sessionResources = Object.assign([bridge, init, packages, stop], { destroy: vi.fn(), }); + + return { bridge, init, packages, sessionResources, stop }; +} + +function createPyodideStub( + sessions = Array.from({ length: 4 }, () => createSessionResources()), +) { + const [first] = sessions; + if (!first) { + throw new Error("At least one session is required"); + } const loadPackagesFromImports = vi.fn().mockResolvedValue(undefined); const pyodide = { - runPython: vi.fn(() => sessionResources), + runPython: vi + .fn() + .mockImplementation(() => sessions.shift()?.sessionResources), loadPackagesFromImports, loadedPackages: {}, runPythonAsync: vi.fn(), } as unknown as PyodideInterface; return { - init, + ...first, loadPackagesFromImports, - packages, pyodide, - sessionResources, - stop, }; } @@ -93,4 +106,50 @@ describe("WASM controller session lifecycle", () => { expect(loadedSources[1]).toContain("current_dependency"); expect(loadedSources.join("\n")).not.toContain("superseded_dependency"); }); + + it("stops every session created for a multi-app page", async () => { + const first = createSessionResources(); + const second = createSessionResources(); + const { pyodide } = createPyodideStub([first, second]); + const controller = new TestController(); + controller.setPyodide(pyodide); + + await startSession(controller, "first"); + await startSession(controller, "second"); + second.bridge.destroy(); + + expect(first.stop.destroy).not.toHaveBeenCalled(); + + await controller.stopSession(); + + expect(first.stop).toHaveBeenCalledOnce(); + expect(second.stop).toHaveBeenCalledOnce(); + expect(first.stop.destroy).toHaveBeenCalledOnce(); + expect(second.stop.destroy).toHaveBeenCalledOnce(); + }); + + it("continues stopping sessions after one stop fails", async () => { + const failure = new Error("stop failed"); + const first = createSessionResources(() => Promise.reject(failure)); + const second = createSessionResources(); + const { pyodide } = createPyodideStub([first, second]); + const controller = new TestController(); + controller.setPyodide(pyodide); + + await startSession(controller, "first"); + await startSession(controller, "second"); + + await expect(controller.stopSession()).rejects.toThrow("stop failed"); + expect(first.stop).toHaveBeenCalledOnce(); + expect(second.stop).toHaveBeenCalledOnce(); + expect(first.stop.destroy).not.toHaveBeenCalled(); + expect(second.stop.destroy).toHaveBeenCalledOnce(); + + first.stop.mockResolvedValueOnce(undefined); + await controller.stopSession(); + + expect(first.stop).toHaveBeenCalledTimes(2); + expect(first.stop.destroy).toHaveBeenCalledOnce(); + expect(second.stop).toHaveBeenCalledOnce(); + }); }); diff --git a/frontend/src/core/wasm/worker/bootstrap.ts b/frontend/src/core/wasm/worker/bootstrap.ts index 2b9318eab24..b14ee170da9 100644 --- a/frontend/src/core/wasm/worker/bootstrap.ts +++ b/frontend/src/core/wasm/worker/bootstrap.ts @@ -29,7 +29,7 @@ export class DefaultWasmController implements WasmController { protected pyodide: PyodideInterface | null = null; private packageLoadQueue = Promise.resolve(); private sessionGeneration = 0; - private stopCurrentSession: PyCallable | undefined; + private activeSessionStops = new Set(); get requirePyodide() { invariant(this.pyodide, "Pyodide not loaded"); @@ -160,23 +160,22 @@ export class DefaultWasmController implements WasmController { async def stop(): nonlocal bridge, session, session_task task = session_task - try: - kernel_task = getattr(session, "kernel_task", None) - if kernel_task is None: - if task is not None: - task.cancel() - else: - kernel_task.stop() + kernel_task = getattr(session, "kernel_task", None) + if kernel_task is None: if task is not None: - try: - await task - except asyncio.CancelledError: - pass - finally: - session_task = None - session = None - bridge = None - gc.collect() + task.cancel() + else: + kernel_task.stop() + if task is not None: + await asyncio.gather(task, return_exceptions=True) + if bridge is not None: + bridge.session = None + task = None + kernel_task = None + session_task = None + session = None + bridge = None + gc.collect() with open("${nbFilename}", "r") as f: packages = session.find_packages(f.read()) @@ -207,8 +206,7 @@ export class DefaultWasmController implements WasmController { } finally { packagesProxy.destroy(); } - this.stopCurrentSession?.destroy(); - this.stopCurrentSession = stopSession; + this.activeSessionStops.add(stopSession); // Fire and forget: // Load notebook dependencies and instantiate the session @@ -239,12 +237,19 @@ export class DefaultWasmController implements WasmController { async stopSession(): Promise { this.sessionGeneration += 1; - const stop = this.stopCurrentSession; - this.stopCurrentSession = undefined; - try { - await stop?.(); - } finally { - stop?.destroy(); + const stops = [...this.activeSessionStops]; + let firstFailure: unknown; + for (const stop of stops) { + try { + await stop(); + this.activeSessionStops.delete(stop); + stop.destroy(); + } catch (error) { + firstFailure ??= error; + } + } + if (firstFailure) { + throw firstFailure; } } From da792ec9b94c2fe7e5199c40d3e1e117e1f2bc14 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?P=C3=A9ter=20Ferenc=20Gyarmati?= Date: Fri, 17 Jul 2026 13:06:29 +0200 Subject: [PATCH 08/12] feat(islands): preserve session ownership across transitions --- .../src/core/islands/__tests__/bridge.test.ts | 212 +++++++++--------- frontend/src/core/islands/bridge.ts | 65 ++---- frontend/src/core/islands/worker/worker.tsx | 22 +- 3 files changed, 136 insertions(+), 163 deletions(-) diff --git a/frontend/src/core/islands/__tests__/bridge.test.ts b/frontend/src/core/islands/__tests__/bridge.test.ts index 9fe9d83efd1..4115c0ee01f 100644 --- a/frontend/src/core/islands/__tests__/bridge.test.ts +++ b/frontend/src/core/islands/__tests__/bridge.test.ts @@ -48,6 +48,8 @@ const { mockReplaceSessionRequest, mockStopSessionRequest, mockStartSessionRequest, + mockMessageListeners, + mockStoreSet, mockParseMarimoIslandApps, mockCreateMarimoFile, mockGetMarimoExportContext, @@ -57,6 +59,8 @@ const { mockReplaceSessionRequest: vi.fn(), mockStopSessionRequest: vi.fn(), mockStartSessionRequest: vi.fn(), + mockMessageListeners: new Map void>(), + mockStoreSet: vi.fn(), mockParseMarimoIslandApps: vi.fn<() => TestIslandApp[]>(() => []), mockCreateMarimoFile: vi.fn(), mockGetMarimoExportContext: vi.fn<() => TestExportContext | undefined>( @@ -78,7 +82,11 @@ vi.mock("@/core/wasm/rpc", () => ({ consumerReady: vi.fn(), }, }, - addMessageListener: vi.fn(), + addMessageListener: vi.fn( + (name: string, listener: (payload: never) => void) => { + mockMessageListeners.set(name, listener); + }, + ), }), })); @@ -103,10 +111,11 @@ vi.mock("@/core/meta/globals", () => ({ })); // Mock the jotai store -vi.mock("@/core/state/jotai", () => ({ +vi.mock("@/core/state/jotai", async (importOriginal) => ({ + ...(await importOriginal()), store: { get: vi.fn(), - set: vi.fn(), + set: mockStoreSet, }, })); @@ -121,18 +130,40 @@ describe("IslandsPyodideBridge", () => { mockParseMarimoIslandApps.mockReturnValue([]); mockCreateMarimoFile.mockReset(); mockGetMarimoExportContext.mockReturnValue(undefined); - bridge = new IslandsPyodideBridge({ autoStartSessions: false }); + mockMessageListeners.clear(); + bridge = new IslandsPyodideBridge(); }); - function mockSingleApp(file = "generated app 1") { - mockParseMarimoIslandApps.mockReturnValue([ - { - id: "app-1", - cells: [{ code: "x = 1", idx: 0, output: "
1
" }], - }, - ]); + function signalWorkerReady() { + const listener = mockMessageListeners.get("ready"); + if (!listener) { + throw new Error("Missing worker ready listener"); + } + listener({} as never); + } + + function app(id = "app-1"): TestIslandApp { + return { + id, + cells: [{ code: "x = 1", idx: 0, output: "
1
" }], + }; + } + + function setSingleApp(file = "generated app 1", appId = "app-1") { + mockParseMarimoIslandApps.mockReturnValue([app(appId)]); mockCreateMarimoFile.mockReturnValue(file); - bridge.workerReady.resolve(); + } + + function mockSingleApp(file = "generated app 1", appId = "app-1") { + setSingleApp(file, appId); + signalWorkerReady(); + } + + function sendSlider() { + return bridge.sendComponentValues({ + objectIds: [uiElementId("slider-1")], + values: [2], + }); } async function initializeSingleApp() { @@ -140,23 +171,17 @@ describe("IslandsPyodideBridge", () => { await bridge.initializeApps(); } - describe("startSessionsForAllApps", () => { + describe("app initialization", () => { it("should prefer trusted export notebook code when there is exactly one reactive app", async () => { - mockParseMarimoIslandApps.mockReturnValue([ - { - id: "app-1", - cells: [{ code: "x = 1", idx: 0, output: "
1
" }], - }, - ]); + mockParseMarimoIslandApps.mockReturnValue([app()]); mockGetMarimoExportContext.mockReturnValue({ trusted: true, notebookCode: "import marimo\napp = marimo.App()\n@app.cell\ndef __():\n x = 1\n return", }); - await ( - bridge as unknown as { startSessionsForAllApps(): Promise } - ).startSessionsForAllApps(); + signalWorkerReady(); + await bridge.initializeApps(); expect(mockCreateMarimoFile).not.toHaveBeenCalled(); expect(mockStartSessionRequest).toHaveBeenCalledWith({ @@ -167,11 +192,7 @@ describe("IslandsPyodideBridge", () => { }); it("should ignore trusted export notebook code for a payload-backed app", async () => { - const payloadApp = { - id: "app-1", - payloadBacked: true, - cells: [{ code: "x = 1", idx: 0, output: "
1
" }], - }; + const payloadApp = { ...app(), payloadBacked: true }; mockParseMarimoIslandApps.mockReturnValue([payloadApp]); mockGetMarimoExportContext.mockReturnValue({ trusted: true, @@ -179,9 +200,8 @@ describe("IslandsPyodideBridge", () => { }); mockCreateMarimoFile.mockReturnValue("generated payload app"); - await ( - bridge as unknown as { startSessionsForAllApps(): Promise } - ).startSessionsForAllApps(); + signalWorkerReady(); + await bridge.initializeApps(); expect(mockCreateMarimoFile).toHaveBeenCalledWith(payloadApp); expect(mockStartSessionRequest).toHaveBeenCalledWith({ @@ -192,16 +212,7 @@ describe("IslandsPyodideBridge", () => { }); it("should keep synthesized per-app files for multiple reactive apps even when export context exists", async () => { - mockParseMarimoIslandApps.mockReturnValue([ - { - id: "app-1", - cells: [{ code: "x = 1", idx: 0, output: "
1
" }], - }, - { - id: "app-2", - cells: [{ code: "y = 2", idx: 0, output: "
2
" }], - }, - ]); + mockParseMarimoIslandApps.mockReturnValue([app(), app("app-2")]); mockGetMarimoExportContext.mockReturnValue({ trusted: true, notebookCode: "full notebook should be ignored", @@ -210,9 +221,8 @@ describe("IslandsPyodideBridge", () => { .mockReturnValueOnce("generated app 1") .mockReturnValueOnce("generated app 2"); - await ( - bridge as unknown as { startSessionsForAllApps(): Promise } - ).startSessionsForAllApps(); + signalWorkerReady(); + await bridge.initializeApps(); expect(mockCreateMarimoFile).toHaveBeenCalledTimes(2); expect(mockStartSessionRequest).toHaveBeenNthCalledWith(1, { @@ -226,10 +236,7 @@ describe("IslandsPyodideBridge", () => { sessionGeneration: 2, }); - await bridge.sendComponentValues({ - objectIds: [uiElementId("slider-1")], - values: [2], - }); + await sendSlider(); await bridge.stopSession(); expect(mockReplaceSessionRequest).not.toHaveBeenCalled(); @@ -248,17 +255,10 @@ describe("IslandsPyodideBridge", () => { }); it("should synthesize a file for a single app when no trusted export context is present", async () => { - mockParseMarimoIslandApps.mockReturnValue([ - { - id: "app-1", - cells: [{ code: "x = 1", idx: 0, output: "
1
" }], - }, - ]); - mockCreateMarimoFile.mockReturnValue("generated app 1"); + setSingleApp(); - await ( - bridge as unknown as { startSessionsForAllApps(): Promise } - ).startSessionsForAllApps(); + signalWorkerReady(); + await bridge.initializeApps(); expect(mockCreateMarimoFile).toHaveBeenCalledTimes(1); expect(mockStartSessionRequest).toHaveBeenCalledWith({ @@ -271,24 +271,15 @@ describe("IslandsPyodideBridge", () => { describe("app lifecycle", () => { it("holds controls until the first app session is ready", async () => { - mockParseMarimoIslandApps.mockReturnValue([ - { - id: "app-1", - cells: [{ code: "x = 1", idx: 0, output: "
1
" }], - }, - ]); - mockCreateMarimoFile.mockReturnValue("generated app 1"); + setSingleApp(); const initialization = bridge.initializeApps(); - const control = bridge.sendComponentValues({ - objectIds: [uiElementId("slider-1")], - values: [2], - }); + const control = sendSlider(); await Promise.resolve(); expect(mockBridge).not.toHaveBeenCalled(); - bridge.workerReady.resolve(); + signalWorkerReady(); await Promise.all([initialization, control]); expect(mockBridge).toHaveBeenCalledWith( expect.objectContaining({ @@ -303,10 +294,7 @@ describe("IslandsPyodideBridge", () => { await bridge.initializeApps(); await bridge.initializeApps(); - await bridge.sendComponentValues({ - objectIds: [uiElementId("slider-1")], - values: [2], - }); + await sendSlider(); expect(mockStartSessionRequest).toHaveBeenCalledOnce(); expect(mockStartSessionRequest).toHaveBeenCalledWith({ @@ -330,6 +318,7 @@ describe("IslandsPyodideBridge", () => { code: "generated app 2", sessionGeneration: 2, }); + expect(mockStoreSet).toHaveBeenCalledOnce(); }); it("stops the matching active app", async () => { @@ -344,6 +333,38 @@ describe("IslandsPyodideBridge", () => { appId: "app-1", sessionGeneration: 1, }); + + mockBridge.mockClear(); + const control = sendSlider(); + await Promise.resolve(); + expect(mockBridge).not.toHaveBeenCalled(); + + mockSingleApp("generated app 2"); + await bridge.initializeApps(); + await control; + + expect(mockBridge).toHaveBeenCalledWith( + expect.objectContaining({ sessionGeneration: 2 }), + ); + }); + + it("replaces the active session after stop fails", async () => { + mockSingleApp(); + await bridge.initializeApps(); + mockStopSessionRequest.mockRejectedValueOnce(new Error("stop failed")); + + await expect(bridge.stopSession("app-1")).rejects.toThrow("stop failed"); + + mockSingleApp("generated app 2", "app-2"); + + await bridge.initializeApps(); + + expect(mockStartSessionRequest).toHaveBeenCalledOnce(); + expect(mockReplaceSessionRequest).toHaveBeenCalledWith({ + appId: "app-2", + code: "generated app 2", + sessionGeneration: 2, + }); }); it("keeps controls sent during stop scoped to the stopped app", async () => { @@ -361,17 +382,8 @@ describe("IslandsPyodideBridge", () => { await vi.waitFor(() => expect(mockStopSessionRequest).toHaveBeenCalledOnce(), ); - const control = bridge.sendComponentValues({ - objectIds: [uiElementId("slider-1")], - values: [2], - }); - mockParseMarimoIslandApps.mockReturnValue([ - { - id: "app-2", - cells: [{ code: "x = 2", idx: 0, output: "
2
" }], - }, - ]); - mockCreateMarimoFile.mockReturnValue("generated app 2"); + const control = sendSlider(); + mockSingleApp("generated app 2", "app-2"); const nextInitialization = bridge.initializeApps(); finishStop(); @@ -398,9 +410,9 @@ describe("IslandsPyodideBridge", () => { ); await bridge.initializeApps(); - expect(mockReplaceSessionRequest).toHaveBeenCalledOnce(); - expect(mockStartSessionRequest).toHaveBeenCalledTimes(2); - expect(mockStartSessionRequest).toHaveBeenLastCalledWith({ + expect(mockStartSessionRequest).toHaveBeenCalledOnce(); + expect(mockReplaceSessionRequest).toHaveBeenCalledTimes(2); + expect(mockReplaceSessionRequest).toHaveBeenLastCalledWith({ appId: "app-1", code: "generated app 2", sessionGeneration: 3, @@ -430,17 +442,8 @@ describe("IslandsPyodideBridge", () => { await vi.waitFor(() => expect(mockReplaceSessionRequest).toHaveBeenCalledOnce(), ); - const control = bridge.sendComponentValues({ - objectIds: [uiElementId("slider-1")], - values: [2], - }); - mockParseMarimoIslandApps.mockReturnValue([ - { - id: "app-2", - cells: [{ code: "x = 3", idx: 0, output: "
3
" }], - }, - ]); - mockCreateMarimoFile.mockReturnValue("generated app 3"); + const control = sendSlider(); + mockSingleApp("generated app 3", "app-2"); const thirdInitialization = bridge.initializeApps(); expect(mockBridge).not.toHaveBeenCalled(); @@ -474,20 +477,11 @@ describe("IslandsPyodideBridge", () => { await vi.waitFor(() => expect(mockReplaceSessionRequest).toHaveBeenCalledOnce(), ); - const control = bridge.sendComponentValues({ - objectIds: [uiElementId("slider-1")], - values: [2], - }); + const control = sendSlider(); failReplacement(new Error("replacement failed")); await expect(failedInitialization).rejects.toThrow("replacement failed"); - mockParseMarimoIslandApps.mockReturnValue([ - { - id: "app-2", - cells: [{ code: "x = 3", idx: 0, output: "
3
" }], - }, - ]); - mockCreateMarimoFile.mockReturnValue("generated app 3"); + mockSingleApp("generated app 3", "app-2"); await Promise.all([bridge.initializeApps(), control]); expect(mockBridge).toHaveBeenCalledWith( diff --git a/frontend/src/core/islands/bridge.ts b/frontend/src/core/islands/bridge.ts index f077823809c..2ccc781618c 100644 --- a/frontend/src/core/islands/bridge.ts +++ b/frontend/src/core/islands/bridge.ts @@ -7,6 +7,7 @@ import { throwNotImplemented } from "@/utils/functions"; import type { JsonString } from "@/utils/json/base64"; import { Logger } from "@/utils/Logger"; import { generateUUID } from "@/utils/uuid"; +import { initialNotebookState, notebookAtom } from "../cells/cells"; import type { CommandMessage, NotificationPayload } from "../kernel/messages"; import type { EditRequests, RunRequests } from "../network/types"; import { store as defaultStore } from "../state/jotai"; @@ -35,11 +36,6 @@ export interface IslandsBridgeConfig { * Optional root element for parsing islands (for testing) */ root?: Document | Element; - - /** - * Whether to auto-start sessions on worker ready (default: true) - */ - autoStartSessions?: boolean; } interface AppSession { @@ -57,8 +53,8 @@ interface AppSession { * @example * ```ts * const bridge = new IslandsPyodideBridge(); - * await bridge.initialized; * bridge.consumeMessages(message => console.log(message)); + * await bridge.initializeApps(); * ``` */ export class IslandsPyodideBridge implements RunRequests, EditRequests { @@ -68,19 +64,17 @@ export class IslandsPyodideBridge implements RunRequests, EditRequests { | undefined; private readonly store: typeof defaultStore; private readonly root: Document | Element; - private readonly autoStartSessions: boolean; private session: AppSession | undefined; private sessionReady = new Deferred(); private nextSessionGeneration = 0; private appTransition = Promise.resolve(); + private workerReady = new Deferred(); public initialized = new Deferred(); - public workerReady = new Deferred(); constructor(config: IslandsBridgeConfig = {}) { this.store = config.store || defaultStore; this.root = config.root || document; - this.autoStartSessions = config.autoStartSessions ?? true; try { const factory = config.workerFactory || new DefaultWorkerFactory(); @@ -100,9 +94,6 @@ export class IslandsPyodideBridge implements RunRequests, EditRequests { private setupMessageListeners(): void { this.rpc.addMessageListener("ready", () => { this.workerReady.resolve(); - if (this.autoStartSessions) { - void this.startSessionsForAllApps(); - } }); this.rpc.addMessageListener("initialized", () => { @@ -127,23 +118,16 @@ export class IslandsPyodideBridge implements RunRequests, EditRequests { ); } - /** - * Starts sessions for all apps found in the DOM - */ - private startSessionsForAllApps(): Promise { - return this.enqueueAppTransition(() => this.startApps(false)); - } - async initializeApps(): Promise { await this.enqueueAppTransition(async () => { await this.workerReady.promise; - await this.startApps(true); + await this.startApps(); }); } - private async startApps(replaceSingleApp: boolean): Promise { + private async startApps(): Promise { const apps = parseMarimoIslandApps(this.root); - const managesSingleApp = replaceSingleApp && apps.length === 1; + const managesSingleApp = apps.length === 1; Logger.debug( `Starting sessions for ${apps.length} app(s):`, apps.map((a) => `${a.id} (${a.cells.length} cells)`), @@ -173,6 +157,9 @@ export class IslandsPyodideBridge implements RunRequests, EditRequests { }; const previousSession = this.session; const replacesSession = managesSingleApp && previousSession !== undefined; + if (replacesSession) { + this.store.set(notebookAtom, initialNotebookState()); + } if (managesSingleApp || !previousSession) { if (managesSingleApp && this.sessionReady.status !== "pending") { this.sessionReady = new Deferred(); @@ -192,7 +179,7 @@ export class IslandsPyodideBridge implements RunRequests, EditRequests { } } catch (error) { if (this.session?.sessionGeneration === request.sessionGeneration) { - this.session = managesSingleApp ? undefined : previousSession; + this.session = previousSession; } Logger.error(`Failed to start session for app ${app.id}:`, error); if (managesSingleApp) { @@ -208,15 +195,13 @@ export class IslandsPyodideBridge implements RunRequests, EditRequests { if (session?.code === undefined || (appId && session.appId !== appId)) { return; } - try { - await this.rpc.proxy.request.stopSession( - appId === undefined - ? { sessionGeneration: session.sessionGeneration } - : { appId, sessionGeneration: session.sessionGeneration }, - ); - } finally { - this.session = undefined; - } + await this.rpc.proxy.request.stopSession({ + appId: session.appId, + sessionGeneration: session.sessionGeneration, + }); + this.session = undefined; + this.sessionReady = new Deferred(); + this.store.set(notebookAtom, initialNotebookState()); }); } @@ -291,9 +276,9 @@ export class IslandsPyodideBridge implements RunRequests, EditRequests { sendRun: EditRequests["sendRun"] = async (request): Promise => { const session = await this.getActiveSession(); await this.rpc.proxy.request.loadPackages({ - appId: session?.appId, + appId: session.appId, code: request.codes.join("\n"), - sessionGeneration: session?.sessionGeneration, + sessionGeneration: session.sessionGeneration, }); await this.putControlRequest( { @@ -378,12 +363,10 @@ export class IslandsPyodideBridge implements RunRequests, EditRequests { ): Promise { session ??= await this.getActiveSession(); await this.rpc.proxy.request.bridge({ + appId: session.appId, functionName: "put_control_request", payload: operation, - ...(session && { - appId: session.appId, - sessionGeneration: session.sessionGeneration, - }), + sessionGeneration: session.sessionGeneration, }); } @@ -392,7 +375,7 @@ export class IslandsPyodideBridge implements RunRequests, EditRequests { const session = this.session; const ready = this.sessionReady; await transition; - return this.session ?? session ?? ready.promise; + return session ?? ready.promise; } /** @@ -411,9 +394,7 @@ let globalBridgeInstance: IslandsPyodideBridge | null = null; export function getGlobalBridge(): IslandsPyodideBridge { if (!globalBridgeInstance) { - globalBridgeInstance = new IslandsPyodideBridge({ - autoStartSessions: false, - }); + globalBridgeInstance = new IslandsPyodideBridge(); } return globalBridgeInstance; } diff --git a/frontend/src/core/islands/worker/worker.tsx b/frontend/src/core/islands/worker/worker.tsx index 83bd71d8837..f17668c6e60 100644 --- a/frontend/src/core/islands/worker/worker.tsx +++ b/frontend/src/core/islands/worker/worker.tsx @@ -80,10 +80,9 @@ async function startSession( ...notebook, onMessage: messageBuffer.push, }); - if (activeSession && !replace) { + if (activeSession) { (nextBridge as unknown as PyProxy).destroy(); } else { - (activeSession?.bridge as unknown as PyProxy | undefined)?.destroy(); activeSession = { appId: opts.appId, bridge: nextBridge, @@ -101,9 +100,9 @@ async function startSession( async function stopActiveSession(): Promise { const session = activeSession; + await self.controller.stopSession(); activeSession = undefined; (session?.bridge as unknown as PyProxy | undefined)?.destroy(); - await self.controller.stopSession(); } function enqueueSession(operation: () => Promise): Promise { @@ -128,10 +127,10 @@ const requestHandler = createRPCRequestHandler({ await enqueueSession(() => startSession(opts, true)); }, - stopSession: async (opts: { appId?: string; sessionGeneration: number }) => { + stopSession: async (opts: { appId: string; sessionGeneration: number }) => { await enqueueSession(async () => { if ( - (opts.appId && activeSession?.appId !== opts.appId) || + activeSession?.appId !== opts.appId || activeSession?.sessionGeneration !== opts.sessionGeneration ) { return; @@ -145,9 +144,9 @@ const requestHandler = createRPCRequestHandler({ * Load packages */ loadPackages: async (opts: { - appId?: string; + appId: string; code: string; - sessionGeneration?: number; + sessionGeneration: number; }) => { await pyodideReadyPromise; // Make sure loading is done await enqueueSession(async () => { @@ -180,10 +179,10 @@ const requestHandler = createRPCRequestHandler({ * Call a function on the bridge */ bridge: async (opts: { - appId?: string; + appId: string; functionName: keyof RawBridge; payload: {} | undefined | null; - sessionGeneration?: number; + sessionGeneration: number; }) => { await pyodideReadyPromise; // Make sure loading is done @@ -216,13 +215,12 @@ const requestHandler = createRPCRequestHandler({ }); function requireActiveBridge(opts: { - appId?: string; - sessionGeneration?: number; + appId: string; + sessionGeneration: number; }): SerializedBridge { const session = activeSession; if ( !session || - !opts.appId || opts.appId !== session.appId || opts.sessionGeneration !== session.sessionGeneration ) { From 03c6f9697ae18758400ff9835332f840d3f816ca Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?P=C3=A9ter=20Ferenc=20Gyarmati?= Date: Fri, 17 Jul 2026 13:06:54 +0200 Subject: [PATCH 09/12] feat(islands): refresh retained page bindings --- .../src/core/islands/__tests__/parse.test.ts | 42 +++- frontend/src/core/islands/bootstrap.ts | 7 +- .../__tests__/web-components.test.tsx | 214 ++++++++++++++++++ .../islands/components/web-components.tsx | 89 ++++++-- frontend/src/core/islands/constants.ts | 1 + frontend/src/core/islands/main.ts | 10 +- frontend/src/core/islands/parse.ts | 17 +- 7 files changed, 351 insertions(+), 29 deletions(-) create mode 100644 frontend/src/core/islands/components/__tests__/web-components.test.tsx diff --git a/frontend/src/core/islands/__tests__/parse.test.ts b/frontend/src/core/islands/__tests__/parse.test.ts index 18f6f3197d0..96ecef186e0 100644 --- a/frontend/src/core/islands/__tests__/parse.test.ts +++ b/frontend/src/core/islands/__tests__/parse.test.ts @@ -54,6 +54,7 @@ function appendPayload( script.type = ISLANDS_JSON_SCRIPT_TYPE; script.textContent = JSON.stringify(payload); root.appendChild(script); + return script; } describe("createMarimoFile", () => { @@ -330,6 +331,9 @@ describe("parseIslandElement", () => { code: 'print("retained")', output: "
retained output
", }); + + retainIslandSource(element, { code: "", output: "
static
" }); + expect(parseIslandElement(element)).toBeNull(); }); it("prefers new live source when an island element is reused", () => { @@ -687,28 +691,48 @@ describe("parseMarimoIslandApps", () => { expect(island.getAttribute(ISLAND_DATA_ATTRIBUTES.CELL_IDX)).toBe("0"); }); - it("keeps a payload anchor stable during duplicate initialization", () => { + it("refreshes a reused payload anchor when its content changes", () => { const island = createMockIslandElement({ appId: "app1", cellId: "cell-1", - code: "dom_code = True", + code: 'print("dom")', innerHTML: "
dom output
", }); + island.insertAdjacentHTML( + "beforeend", + '
', + ); island.setAttribute(ISLAND_DATA_ATTRIBUTES.REACTIVE, "true"); container.appendChild(island); - appendPayload(container, { + const script = appendPayload(container, { schemaVersion: 1, appId: "app1", cells: [createPayloadCell()], }); + parseMarimoIslandApps(container); + + script.textContent = JSON.stringify({ + schemaVersion: 1, + appId: "app1", + cells: [ + createPayloadCell({ + code: 'print("updated")', + outputHtml: "
updated output
", + }), + ], + }); - const first = parseMarimoIslandApps(container); - island.replaceChildren(); - const second = parseMarimoIslandApps(container); + parseMarimoIslandApps(container); - expect(second).toEqual(first); - expect(island.childNodes).toHaveLength(0); - expect(island.getAttribute(ISLAND_DATA_ATTRIBUTES.CELL_IDX)).toBe("0"); + expect(extractIslandCodeFromEmbed(island)).toBe('print("updated")'); + expect(island.querySelector(ISLAND_TAG_NAMES.CELL_OUTPUT)?.innerHTML).toBe( + "
updated output
", + ); + expect( + island + .querySelector(ISLAND_TAG_NAMES.CODE_EDITOR) + ?.getAttribute("data-initial-value"), + ).toBe(JSON.stringify('print("updated")')); }); it("should parse Python-generated island payload snapshots", () => { diff --git a/frontend/src/core/islands/bootstrap.ts b/frontend/src/core/islands/bootstrap.ts index c60dbee0d40..f132080665e 100644 --- a/frontend/src/core/islands/bootstrap.ts +++ b/frontend/src/core/islands/bootstrap.ts @@ -108,14 +108,17 @@ export async function initializeIslands( // Loading indicator: dim islands while Pyodide initializes store.sub(shouldShowIslandsWarningIndicatorAtom, () => { const showing = store.get(shouldShowIslandsWarningIndicatorAtom); + const currentIslands = root.querySelectorAll( + ISLAND_TAG_NAMES.ISLAND, + ); if (showing) { toastIslandsLoading(); - for (const island of islands) { + for (const island of currentIslands) { island.style.setProperty("opacity", "0.5"); } } else { dismissIslandsLoadingToast(); - for (const island of islands) { + for (const island of currentIslands) { island.style.removeProperty("opacity"); } } diff --git a/frontend/src/core/islands/components/__tests__/web-components.test.tsx b/frontend/src/core/islands/components/__tests__/web-components.test.tsx new file mode 100644 index 00000000000..68895dd3fb0 --- /dev/null +++ b/frontend/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 = ` + +
initial output
 + + + `; + await actAndFlush(() => document.body.append(container)); + + const island = container.querySelector( + 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 = ` + +
old output
 + + + `; + const script = document.createElement("script"); + script.type = ISLANDS_JSON_SCRIPT_TYPE; + script.textContent = payloadSource({ + code: 'print("unchanged")', + outputHtml: "
old output
", + }); + container.append(script); + await actAndFlush(() => document.body.append(container)); + + await actAndFlush(() => { + parseMarimoIslandApps(container); + script.textContent = payloadSource({ + code: 'print("unchanged")', + outputHtml: "
updated output
", + }); + parseMarimoIslandApps(container); + }); + + const island = container.querySelector( + 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: "
next app output
", + }); + 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: `
${output}
`, + 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 { + 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; +} diff --git a/frontend/src/core/islands/components/web-components.tsx b/frontend/src/core/islands/components/web-components.tsx index 88a9d500fca..f6c5cbd942e 100644 --- a/frontend/src/core/islands/components/web-components.tsx +++ b/frontend/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,6 +16,7 @@ import { store } from "../../state/jotai"; import { ISLAND_CSS_CLASSES, ISLAND_DATA_ATTRIBUTES, + ISLAND_SOURCE_CHANGED_EVENT, ISLAND_TAG_NAMES, } from "../constants"; import { extractIslandCodeFromEmbed, retainIslandSource } from "../parse"; @@ -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,6 @@ 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. @@ -153,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 === "" || trimmedHtml === ""; @@ -169,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( @@ -243,6 +297,13 @@ export class MarimoIslandElement extends HTMLElement { * Cleanup when element is removed from DOM */ disconnectedCallback(): void { + this.connectionGeneration += 1; + 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 diff --git a/frontend/src/core/islands/constants.ts b/frontend/src/core/islands/constants.ts index a54ca97d1de..9167b2d9424 100644 --- a/frontend/src/core/islands/constants.ts +++ b/frontend/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 diff --git a/frontend/src/core/islands/main.ts b/frontend/src/core/islands/main.ts index 3dd9c57c828..e2721cad753 100644 --- a/frontend/src/core/islands/main.ts +++ b/frontend/src/core/islands/main.ts @@ -14,7 +14,7 @@ import "iconify-icon"; import { Logger } from "@/utils/Logger"; import { initializeIslands } from "./bootstrap"; import { getGlobalBridge } from "./bridge"; -import { ISLAND_TAG_NAMES } from "./constants"; +import { ISLAND_CSS_CLASSES, ISLAND_TAG_NAMES } from "./constants"; import { parseMarimoIslandApps } from "./parse"; const bridge = getGlobalBridge(); @@ -32,9 +32,15 @@ export function canReplaceApp(): boolean { * Mounts marimo custom elements and starts the apps in the current document. */ export async function initialize(): Promise { - if (!document.querySelector(ISLAND_TAG_NAMES.ISLAND)) { + const islands = document.querySelectorAll( + 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; diff --git a/frontend/src/core/islands/parse.ts b/frontend/src/core/islands/parse.ts index b163627cd0e..d6237caa16d 100644 --- a/frontend/src/core/islands/parse.ts +++ b/frontend/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"; @@ -147,6 +148,7 @@ export function parseIslandElementsIntoApps( // Add data-cell-idx attribute to the island element if (materialize) { embed.setAttribute(ISLAND_DATA_ATTRIBUTES.CELL_IDX, idx.toString()); + embed.dispatchEvent(new Event(ISLAND_SOURCE_CHANGED_EVENT)); } } @@ -181,7 +183,7 @@ function parsePayloadBackedApps({ } consumedEmbeds.add(embed); matchedPayloadCells.set(cell, embed); - if (materialize && !retainedPayloadCellIds.has(embed)) { + if (materialize) { materializeIslandPayload(embed, cell); } hasMatchedIsland = true; @@ -244,6 +246,7 @@ function parsePayloadBackedApps({ embed.removeAttribute(ISLAND_DATA_ATTRIBUTES.CELL_ID); embed.removeAttribute(ISLAND_DATA_ATTRIBUTES.CELL_IDX); embed.setAttribute(ISLAND_DATA_ATTRIBUTES.REACTIVE, "false"); + consumedEmbeds.add(embed); } } } @@ -253,6 +256,12 @@ function parsePayloadBackedApps({ return !appId || !payloadAppIds.has(appId); }); + if (materialize) { + for (const embed of consumedEmbeds) { + embed.dispatchEvent(new Event(ISLAND_SOURCE_CHANGED_EVENT)); + } + } + return [ ...apps.values(), ...parseIslandElementsIntoApps(domOnlyEmbeds, materialize), @@ -334,7 +343,11 @@ export function retainIslandSource( embed: HTMLElement, source: { output: string; code: string }, ): void { - retainedIslandSources.set(embed, source); + if (source.code) { + retainedIslandSources.set(embed, source); + } else { + retainedIslandSources.delete(embed); + } } export function parseIslandElement( From fa1a3fcedae5d04f157c0dca4ff64ffcc0ebf4ac Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?P=C3=A9ter=20Ferenc=20Gyarmati?= Date: Fri, 17 Jul 2026 13:47:55 +0200 Subject: [PATCH 10/12] fix(wasm): propagate falsy teardown failures --- .../worker/__tests__/controller.test.ts | 18 ++++++++++++++++++ frontend/src/core/wasm/worker/bootstrap.ts | 8 ++++++-- 2 files changed, 24 insertions(+), 2 deletions(-) diff --git a/frontend/src/core/islands/worker/__tests__/controller.test.ts b/frontend/src/core/islands/worker/__tests__/controller.test.ts index b86a8a646e0..1fc00ec5184 100644 --- a/frontend/src/core/islands/worker/__tests__/controller.test.ts +++ b/frontend/src/core/islands/worker/__tests__/controller.test.ts @@ -152,4 +152,22 @@ describe("WASM controller session lifecycle", () => { expect(first.stop.destroy).toHaveBeenCalledOnce(); expect(second.stop).toHaveBeenCalledOnce(); }); + + it("reports and retries a falsy stop failure", async () => { + const session = createSessionResources(() => Promise.reject(false)); + const { pyodide } = createPyodideStub([session]); + const controller = new TestController(); + controller.setPyodide(pyodide); + + await startSession(controller, "current"); + + await expect(controller.stopSession()).rejects.toBe(false); + expect(session.stop.destroy).not.toHaveBeenCalled(); + + session.stop.mockResolvedValueOnce(undefined); + await controller.stopSession(); + + expect(session.stop).toHaveBeenCalledTimes(2); + expect(session.stop.destroy).toHaveBeenCalledOnce(); + }); }); diff --git a/frontend/src/core/wasm/worker/bootstrap.ts b/frontend/src/core/wasm/worker/bootstrap.ts index b14ee170da9..b6c997ae72d 100644 --- a/frontend/src/core/wasm/worker/bootstrap.ts +++ b/frontend/src/core/wasm/worker/bootstrap.ts @@ -238,6 +238,7 @@ export class DefaultWasmController implements WasmController { async stopSession(): Promise { this.sessionGeneration += 1; const stops = [...this.activeSessionStops]; + let hasFailure = false; let firstFailure: unknown; for (const stop of stops) { try { @@ -245,10 +246,13 @@ export class DefaultWasmController implements WasmController { this.activeSessionStops.delete(stop); stop.destroy(); } catch (error) { - firstFailure ??= error; + if (!hasFailure) { + hasFailure = true; + firstFailure = error; + } } } - if (firstFailure) { + if (hasFailure) { throw firstFailure; } } From 4e220f0c6f74aab9ae1229ce69908e4dd0889126 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?P=C3=A9ter=20Ferenc=20Gyarmati?= Date: Fri, 17 Jul 2026 17:23:36 +0200 Subject: [PATCH 11/12] simplify connection lifecycle guard --- frontend/src/core/islands/components/web-components.tsx | 1 - 1 file changed, 1 deletion(-) diff --git a/frontend/src/core/islands/components/web-components.tsx b/frontend/src/core/islands/components/web-components.tsx index f6c5cbd942e..121941129ae 100644 --- a/frontend/src/core/islands/components/web-components.tsx +++ b/frontend/src/core/islands/components/web-components.tsx @@ -297,7 +297,6 @@ export class MarimoIslandElement extends HTMLElement { * Cleanup when element is removed from DOM */ disconnectedCallback(): void { - this.connectionGeneration += 1; this.removeEventListener( ISLAND_SOURCE_CHANGED_EVENT, this.handleSourceChanged, From 4eec1e6a8f97f533f83ecac2cd11f485f72f8481 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?P=C3=A9ter=20Ferenc=20Gyarmati?= Date: Fri, 17 Jul 2026 17:23:37 +0200 Subject: [PATCH 12/12] use typed tuple --- frontend/src/core/wasm/worker/bootstrap.ts | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/frontend/src/core/wasm/worker/bootstrap.ts b/frontend/src/core/wasm/worker/bootstrap.ts index b6c997ae72d..91ae31d88eb 100644 --- a/frontend/src/core/wasm/worker/bootstrap.ts +++ b/frontend/src/core/wasm/worker/bootstrap.ts @@ -13,7 +13,12 @@ import type { SerializedBridge, WasmController } from "./types"; import { shouldLoadDuckDBPackages } from "../utils"; const MAKE_SNAPSHOT = false; -type SessionResources = [PyProxy, PyCallable, PyProxy, PyCallable]; +type SessionResources = [ + bridge: PyProxy, + init: PyCallable, + packages: PyProxy, + stop: PyCallable, +]; // This class initializes the wasm environment // We would like this initialization to be parallelizable