diff --git a/frontend/src/core/islands/__tests__/bridge.test.ts b/frontend/src/core/islands/__tests__/bridge.test.ts index cac5fa2e9ca..4115c0ee01f 100644 --- a/frontend/src/core/islands/__tests__/bridge.test.ts +++ b/frontend/src/core/islands/__tests__/bridge.test.ts @@ -45,14 +45,22 @@ vi.stubGlobal("URL", MockURL); const { mockBridge, mockLoadPackages, + mockReplaceSessionRequest, + mockStopSessionRequest, mockStartSessionRequest, + mockMessageListeners, + mockStoreSet, mockParseMarimoIslandApps, mockCreateMarimoFile, mockGetMarimoExportContext, } = vi.hoisted(() => ({ mockBridge: vi.fn(), mockLoadPackages: vi.fn(), + 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>( @@ -66,13 +74,19 @@ vi.mock("@/core/wasm/rpc", () => ({ request: { bridge: mockBridge, loadPackages: mockLoadPackages, + replaceSession: mockReplaceSessionRequest, startSession: mockStartSessionRequest, + stopSession: mockStopSessionRequest, }, send: { consumerReady: vi.fn(), }, }, - addMessageListener: vi.fn(), + addMessageListener: vi.fn( + (name: string, listener: (payload: never) => void) => { + mockMessageListeners.set(name, listener); + }, + ), }), })); @@ -97,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, }, })); @@ -115,40 +130,69 @@ describe("IslandsPyodideBridge", () => { mockParseMarimoIslandApps.mockReturnValue([]); mockCreateMarimoFile.mockReset(); mockGetMarimoExportContext.mockReturnValue(undefined); - bridge = new IslandsPyodideBridge({ autoStartSessions: false }); + mockMessageListeners.clear(); + bridge = new IslandsPyodideBridge(); }); - describe("startSessionsForAllApps", () => { + 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); + } + + 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() { + mockSingleApp(); + await bridge.initializeApps(); + } + + 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({ appId: "app-1", code: "import marimo\napp = marimo.App()\n@app.cell\ndef __():\n x = 1\n return", + sessionGeneration: 1, }); }); 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, @@ -156,28 +200,19 @@ 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({ appId: "app-1", code: "generated payload app", + sessionGeneration: 1, }); }); 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", @@ -186,43 +221,281 @@ 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, { appId: "app-1", code: "generated app 1", + sessionGeneration: 1, }); expect(mockStartSessionRequest).toHaveBeenNthCalledWith(2, { appId: "app-2", code: "generated app 2", + sessionGeneration: 2, + }); + + await sendSlider(); + 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, }); }); 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({ appId: "app-1", code: "generated app 1", + sessionGeneration: 1, + }); + }); + }); + + describe("app lifecycle", () => { + it("holds controls until the first app session is ready", async () => { + setSingleApp(); + + const initialization = bridge.initializeApps(); + const control = sendSlider(); + await Promise.resolve(); + + expect(mockBridge).not.toHaveBeenCalled(); + + signalWorkerReady(); + 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(); + await bridge.initializeApps(); + await sendSlider(); + + 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", + sessionGeneration: 1, + }), + ); + mockCreateMarimoFile.mockReturnValue("generated app 2"); + await bridge.initializeApps(); + + expect(mockReplaceSessionRequest).toHaveBeenCalledOnce(); + expect(mockReplaceSessionRequest).toHaveBeenCalledWith({ + appId: "app-1", + code: "generated app 2", + sessionGeneration: 2, + }); + expect(mockStoreSet).toHaveBeenCalledOnce(); + }); + + 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, + }); + + 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 () => { + 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 = sendSlider(); + mockSingleApp("generated app 2", "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(); + 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(mockStartSessionRequest).toHaveBeenCalledOnce(); + expect(mockReplaceSessionRequest).toHaveBeenCalledTimes(2); + expect(mockReplaceSessionRequest).toHaveBeenLastCalledWith({ + appId: "app-1", + code: "generated app 2", + sessionGeneration: 3, }); }); + + 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 = sendSlider(); + mockSingleApp("generated app 3", "app-2"); + 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 failedInitialization = bridge.initializeApps(); + await vi.waitFor(() => + expect(mockReplaceSessionRequest).toHaveBeenCalledOnce(), + ); + const control = sendSlider(); + failReplacement(new Error("replacement failed")); + await expect(failedInitialization).rejects.toThrow("replacement failed"); + + mockSingleApp("generated app 3", "app-2"); + await Promise.all([bridge.initializeApps(), control]); + + expect(mockBridge).toHaveBeenCalledWith( + expect.objectContaining({ + appId: "app-1", + sessionGeneration: 2, + }), + ); + }); }); describe("sendComponentValues", () => { + beforeEach(initializeSingleApp); + it("should include type field and token in control request", async () => { const request = { objectIds: [uiElementId("Hbol-0")], @@ -232,6 +505,7 @@ describe("IslandsPyodideBridge", () => { await bridge.sendComponentValues(request); expect(mockBridge).toHaveBeenCalledWith({ + appId: "app-1", functionName: "put_control_request", payload: { type: "update-ui-element", @@ -239,6 +513,7 @@ describe("IslandsPyodideBridge", () => { values: [58], token: "test-uuid-12345", }, + sessionGeneration: 1, }); }); @@ -251,17 +526,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"), @@ -273,6 +552,7 @@ describe("IslandsPyodideBridge", () => { await bridge.sendFunctionRequest(request); expect(mockBridge).toHaveBeenCalledWith({ + appId: "app-1", functionName: "put_control_request", payload: { type: "invoke-function", @@ -281,11 +561,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")], @@ -295,12 +578,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, }); }); @@ -313,7 +598,11 @@ describe("IslandsPyodideBridge", () => { await bridge.sendRun(request); // Verify loadPackages was called with joined codes - expect(mockLoadPackages).toHaveBeenCalledWith("import pandas"); + expect(mockLoadPackages).toHaveBeenCalledWith({ + appId: "app-1", + code: "import pandas", + sessionGeneration: 1, + }); // Verify order: loadPackages should be called before bridge const loadPackagesCallOrder = @@ -324,6 +613,8 @@ describe("IslandsPyodideBridge", () => { }); describe("sendModelValue", () => { + beforeEach(initializeSingleApp); + it("should include type field in control request", async () => { const request = { modelId: widgetModelId("widget-1"), @@ -338,6 +629,7 @@ describe("IslandsPyodideBridge", () => { await bridge.sendModelValue(request); expect(mockBridge).toHaveBeenCalledWith({ + appId: "app-1", functionName: "put_control_request", payload: { type: "model", @@ -349,11 +641,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/__tests__/parse.test.ts b/frontend/src/core/islands/__tests__/parse.test.ts index 593170e28fb..96ecef186e0 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"; @@ -53,6 +54,7 @@ function appendPayload( script.type = ISLANDS_JSON_SCRIPT_TYPE; script.textContent = JSON.stringify(payload); root.appendChild(script); + return script; } describe("createMarimoFile", () => { @@ -317,6 +319,39 @@ 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
", + }); + + retainIslandSource(element, { code: "", output: "
static
" }); + expect(parseIslandElement(element)).toBeNull(); + }); + + 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", () => { @@ -630,6 +665,76 @@ 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("refreshes a reused payload anchor when its content changes", () => { + const island = createMockIslandElement({ + appId: "app1", + cellId: "cell-1", + code: 'print("dom")', + innerHTML: "
dom output
", + }); + island.insertAdjacentHTML( + "beforeend", + '
', + ); + island.setAttribute(ISLAND_DATA_ATTRIBUTES.REACTIVE, "true"); + container.appendChild(island); + 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
", + }), + ], + }); + + parseMarimoIslandApps(container); + + 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", () => { const html = readFileSync( new URL( diff --git a/frontend/src/core/islands/bootstrap.ts b/frontend/src/core/islands/bootstrap.ts index be1eaa0086e..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"); } } @@ -208,7 +211,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/bridge.ts b/frontend/src/core/islands/bridge.ts index e3ed684d3f0..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,12 @@ 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 { + appId: string; + code?: string; + sessionGeneration: number; } /** @@ -51,8 +53,8 @@ export interface IslandsBridgeConfig { * @example * ```ts * const bridge = new IslandsPyodideBridge(); - * await bridge.initialized; * bridge.consumeMessages(message => console.log(message)); + * await bridge.initializeApps(); * ``` */ export class IslandsPyodideBridge implements RunRequests, EditRequests { @@ -62,14 +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(); 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(); @@ -88,9 +93,7 @@ export class IslandsPyodideBridge implements RunRequests, EditRequests { */ private setupMessageListeners(): void { this.rpc.addMessageListener("ready", () => { - if (this.autoStartSessions) { - this.startSessionsForAllApps(); - } + this.workerReady.resolve(); }); this.rpc.addMessageListener("initialized", () => { @@ -115,11 +118,16 @@ export class IslandsPyodideBridge implements RunRequests, EditRequests { ); } - /** - * Starts sessions for all apps found in the DOM - */ - private startSessionsForAllApps(): void { + async initializeApps(): Promise { + await this.enqueueAppTransition(async () => { + await this.workerReady.promise; + await this.startApps(); + }); + } + + private async startApps(): Promise { const apps = parseMarimoIslandApps(this.root); + const managesSingleApp = apps.length === 1; Logger.debug( `Starting sessions for ${apps.length} app(s):`, apps.map((a) => `${a.id} (${a.cells.length} cells)`), @@ -134,20 +142,83 @@ export class IslandsPyodideBridge implements RunRequests, EditRequests { const notebookCode = exportContext?.notebookCode; for (const app of apps) { const file = notebookCode || createMarimoFile(app); + if ( + managesSingleApp && + 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; + 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(); + } + this.session = { + ...request, + code: managesSingleApp ? file : undefined, + }; + } + const operation = replacesSession + ? this.rpc.proxy.request.replaceSession(request) + : 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 = previousSession; + } Logger.error(`Failed to start session for app ${app.id}:`, error); - }); + if (managesSingleApp) { + throw error; + } + } } } + async stopSession(appId?: string): Promise { + await this.enqueueAppTransition(async () => { + const session = this.session; + if (session?.code === undefined || (appId && session.appId !== appId)) { + return; + } + await this.rpc.proxy.request.stopSession({ + appId: session.appId, + sessionGeneration: session.sessionGeneration, + }); + this.session = undefined; + this.sessionReady = new Deferred(); + this.store.set(notebookAtom, initialNotebookState()); + }); + } + + 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 +274,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 = await this.getActiveSession(); + 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,13 +357,27 @@ 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?: AppSession, + ): Promise { + session ??= await this.getActiveSession(); await this.rpc.proxy.request.bridge({ + appId: session.appId, functionName: "put_control_request", payload: operation, + sessionGeneration: session.sessionGeneration, }); } + private async getActiveSession(): Promise { + const transition = this.appTransition; + const session = this.session; + const ready = this.sessionReady; + await transition; + return session ?? ready.promise; + } + /** * Cleans up resources (for testing) */ 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 557b5e90027..121941129ae 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,9 +16,10 @@ import { store } from "../../state/jotai"; import { ISLAND_CSS_CLASSES, ISLAND_DATA_ATTRIBUTES, + ISLAND_SOURCE_CHANGED_EVENT, ISLAND_TAG_NAMES, } from "../constants"; -import { extractIslandCodeFromEmbed } from "../parse"; +import { extractIslandCodeFromEmbed, retainIslandSource } from "../parse"; import { MarimoOutputWrapper } from "./output-wrapper"; /** @@ -28,7 +29,6 @@ export interface IslandRenderConfig { html: string; codeCallback: () => string; editor: JSX.Element | null; - cellId: CellId | undefined; } /** @@ -38,7 +38,12 @@ export interface IslandRenderConfig { * functionality like re-running cells and copying code. */ export class MarimoIslandElement extends HTMLElement { + private connectionGeneration = 0; private root?: Root; + private renderConfig?: IslandRenderConfig; + private runtimeAppId?: string; + private runtimeCellId?: CellId; + private unsubscribeCellIds?: () => void; public static readonly tagName = ISLAND_TAG_NAMES.ISLAND; public static readonly outputTagName = ISLAND_TAG_NAMES.CELL_OUTPUT; @@ -72,6 +77,9 @@ export class MarimoIslandElement extends HTMLElement { * Returns undefined for non-reactive islands (they have no corresponding cell). */ get cellId(): CellId | undefined { + if (!this.isReactive) { + return undefined; + } const cellId = this.getAttribute(ISLAND_DATA_ATTRIBUTES.CELL_ID); if (cellId) { return cellId as CellId; @@ -94,11 +102,9 @@ export class MarimoIslandElement extends HTMLElement { /** * Looks up a cell ID from the notebook state by index */ - private getCellIdFromIndex(idx: number): CellId { + private getCellIdFromIndex(idx: number): CellId | undefined { const { cellIds } = store.get(notebookAtom); - const cellId = cellIds.inOrderIds.at(idx); - invariant(cellId, `Missing cell ID at index ${idx}`); - return cellId; + return cellIds.inOrderIds.at(idx); } /** @@ -110,18 +116,62 @@ export class MarimoIslandElement extends HTMLElement { * synchronously from there causes "unmount during render" warnings. */ connectedCallback(): void { - // Capture config synchronously (before children get cleared by createRoot) - const config = this.extractRenderConfig(); + const connectionGeneration = ++this.connectionGeneration; + this.addEventListener( + ISLAND_SOURCE_CHANGED_EVENT, + this.handleSourceChanged, + ); + // Capture config synchronously (before children get cleared by createRoot). + // A reconnected element reuses the source captured on its first mount. + if (this.querySelector(MarimoIslandElement.outputTagName)) { + this.renderConfig = this.extractRenderConfig(); + } + invariant(this.renderConfig, "Missing island render source"); + this.runtimeAppId = this.appId; + this.runtimeCellId = this.hasAttribute(ISLAND_DATA_ATTRIBUTES.CELL_IDX) + ? this.cellId + : undefined; + this.syncCellIdsSubscription(); queueMicrotask(() => { // Guard against disconnect between connectedCallback and microtask - if (!this.isConnected) { + if ( + !this.isConnected || + connectionGeneration !== this.connectionGeneration + ) { return; } this.root = ReactDOM.createRoot(this); - this.renderIsland(config); + this.renderIsland(); }); } + private handleSourceChanged = (): void => { + if (this.querySelector(MarimoIslandElement.outputTagName)) { + this.renderConfig = this.extractRenderConfig(); + this.querySelector(MarimoIslandElement.outputTagName)?.remove(); + this.querySelector(MarimoIslandElement.codeTagName)?.remove(); + } + this.runtimeCellId = + this.isReactive && this.runtimeAppId === this.appId + ? this.cellId + : undefined; + this.runtimeAppId = this.appId; + this.syncCellIdsSubscription(); + this.renderIsland(); + }; + + private syncCellIdsSubscription(): void { + if (this.isReactive) { + this.unsubscribeCellIds ??= store.sub(cellIdsAtom, () => { + this.runtimeCellId = this.cellId; + this.renderIsland(); + }); + return; + } + this.unsubscribeCellIds?.(); + this.unsubscribeCellIds = undefined; + } + /** * Extracts configuration needed for rendering */ @@ -130,7 +180,7 @@ export class MarimoIslandElement extends HTMLElement { const initialOutput = output.innerHTML; const optionalEditor = this.getOptionalEditor(); const code = this.code; - const cellId = this.cellId; + retainIslandSource(this, { code, output: initialOutput }); // Read objectId directly from the DOM before createRoot clears children. // optionalEditor is a wrapper, so its .props don't carry the @@ -152,15 +202,19 @@ export class MarimoIslandElement extends HTMLElement { html: initialOutput, codeCallback, editor: optionalEditor, - cellId, }; } /** * Renders the island with React */ - private renderIsland(config: IslandRenderConfig): void { - const { html, codeCallback, editor, cellId } = config; + private renderIsland(): void { + const config = this.renderConfig; + if (!config) { + return; + } + const { html, codeCallback, editor } = config; + const cellId = this.runtimeCellId; const alwaysShowRun = !!editor; const trimmedHtml = html.trim(); const isEmpty = trimmedHtml === "" || trimmedHtml === ""; @@ -168,6 +222,7 @@ export class MarimoIslandElement extends HTMLElement { // Non-reactive islands have no cell in the kernel — just render static HTML if (!cellId) { + this.removeAttribute("data-status"); this.root?.render( @@ -242,6 +297,12 @@ export class MarimoIslandElement extends HTMLElement { * Cleanup when element is removed from DOM */ disconnectedCallback(): void { + this.removeEventListener( + ISLAND_SOURCE_CHANGED_EVENT, + this.handleSourceChanged, + ); + this.unsubscribeCellIds?.(); + this.unsubscribeCellIds = undefined; const root = this.root; this.root = undefined; // Defer unmount to avoid "unmount during render" race 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 21dd168b8db..e2721cad753 100644 --- a/frontend/src/core/islands/main.ts +++ b/frontend/src/core/islands/main.ts @@ -14,12 +14,44 @@ import "iconify-icon"; import { Logger } from "@/utils/Logger"; import { initializeIslands } from "./bootstrap"; import { getGlobalBridge } from "./bridge"; +import { ISLAND_CSS_CLASSES, ISLAND_TAG_NAMES } from "./constants"; +import { parseMarimoIslandApps } from "./parse"; + +const bridge = getGlobalBridge(); +let bootstrapPromise: Promise | 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; +} /** - * 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 { + 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; + }); + await bootstrapPromise; + 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/parse.ts b/frontend/src/core/islands/parse.ts index 6662b70f665..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"; @@ -74,12 +75,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 +94,12 @@ export function parseMarimoIslandApps( return []; } + const materialize = options.materialize ?? true; if (payloads.length > 0) { - return parsePayloadBackedApps(embeds, payloads); + return parsePayloadBackedApps({ embeds, materialize, payloads }); } - return parseIslandElementsIntoApps(embeds); + return parseIslandElementsIntoApps(embeds, materialize); } /** @@ -103,6 +108,7 @@ export function parseMarimoIslandApps( */ export function parseIslandElementsIntoApps( embeds: HTMLElement[], + materialize = true, ): MarimoIslandApp[] { const apps = new Map(); @@ -140,16 +146,24 @@ export function parseIslandElementsIntoApps( }); // Add data-cell-idx attribute to the island element - embed.setAttribute(ISLAND_DATA_ATTRIBUTES.CELL_IDX, idx.toString()); + if (materialize) { + embed.setAttribute(ISLAND_DATA_ATTRIBUTES.CELL_IDX, idx.toString()); + embed.dispatchEvent(new Event(ISLAND_SOURCE_CHANGED_EVENT)); + } } return [...apps.values()]; } -function parsePayloadBackedApps( - embeds: HTMLElement[], - payloads: MarimoIslandPayload[], -): MarimoIslandApp[] { +function parsePayloadBackedApps({ + embeds, + materialize, + payloads, +}: { + embeds: HTMLElement[]; + materialize: boolean; + payloads: MarimoIslandPayload[]; +}): MarimoIslandApp[] { const apps = new Map(); const matchedPayloadCells = new Map(); const consumedEmbeds = new Set(); @@ -169,7 +183,9 @@ function parsePayloadBackedApps( } consumedEmbeds.add(embed); matchedPayloadCells.set(cell, embed); - materializeIslandPayload(embed, cell); + if (materialize) { + materializeIslandPayload(embed, cell); + } hasMatchedIsland = true; } // Only payloads matched to island anchors can start runtime apps. @@ -215,7 +231,7 @@ function parsePayloadBackedApps( appCell.disabled = true; } app.cells.push(appCell); - if (cell.reactive) { + if (materialize && cell.reactive) { embed?.setAttribute(ISLAND_DATA_ATTRIBUTES.CELL_IDX, idx.toString()); } } @@ -223,12 +239,15 @@ function parsePayloadBackedApps( // A supported payload is the runtime source for its app. Extra same-app DOM // islands are disconnected from runtime binding. - for (const embed of embeds) { - const appId = embed.getAttribute(ISLAND_DATA_ATTRIBUTES.APP_ID); - if (appId && payloadAppIds.has(appId) && !consumedEmbeds.has(embed)) { - embed.removeAttribute(ISLAND_DATA_ATTRIBUTES.CELL_ID); - embed.removeAttribute(ISLAND_DATA_ATTRIBUTES.CELL_IDX); - embed.setAttribute(ISLAND_DATA_ATTRIBUTES.REACTIVE, "false"); + if (materialize) { + for (const embed of embeds) { + const appId = embed.getAttribute(ISLAND_DATA_ATTRIBUTES.APP_ID); + if (appId && payloadAppIds.has(appId) && !consumedEmbeds.has(embed)) { + embed.removeAttribute(ISLAND_DATA_ATTRIBUTES.CELL_ID); + embed.removeAttribute(ISLAND_DATA_ATTRIBUTES.CELL_IDX); + embed.setAttribute(ISLAND_DATA_ATTRIBUTES.REACTIVE, "false"); + consumedEmbeds.add(embed); + } } } @@ -237,7 +256,16 @@ function parsePayloadBackedApps( return !appId || !payloadAppIds.has(appId); }); - return [...apps.values(), ...parseIslandElementsIntoApps(domOnlyEmbeds)]; + if (materialize) { + for (const embed of consumedEmbeds) { + embed.dispatchEvent(new Event(ISLAND_SOURCE_CHANGED_EVENT)); + } + } + + return [ + ...apps.values(), + ...parseIslandElementsIntoApps(domOnlyEmbeds, materialize), + ]; } function findMatchingIsland({ @@ -257,7 +285,8 @@ function findMatchingIsland({ } return ( embed.getAttribute(ISLAND_DATA_ATTRIBUTES.APP_ID) === appId && - embed.getAttribute(ISLAND_DATA_ATTRIBUTES.CELL_ID) === cell.cellId + (embed.getAttribute(ISLAND_DATA_ATTRIBUTES.CELL_ID) ?? + retainedPayloadCellIds.get(embed)) === cell.cellId ); }); } @@ -266,6 +295,7 @@ function materializeIslandPayload( embed: HTMLElement, cell: MarimoIslandPayloadCell, ): void { + retainedPayloadCellIds.set(embed, cell.cellId); embed.setAttribute( ISLAND_DATA_ATTRIBUTES.REACTIVE, JSON.stringify(cell.reactive), @@ -304,6 +334,22 @@ function ensureIslandChild(embed: HTMLElement, tagName: string): HTMLElement { * @param embed - The island HTML element * @returns Cell data or null if invalid */ +const retainedIslandSources = new WeakMap< + HTMLElement, + { output: string; code: string } +>(); + +export function retainIslandSource( + embed: HTMLElement, + source: { output: string; code: string }, +): void { + if (source.code) { + retainedIslandSources.set(embed, source); + } else { + retainedIslandSources.delete(embed); + } +} + export function parseIslandElement( embed: HTMLElement, ): { output: string; code: string } | null { @@ -312,14 +358,15 @@ export function parseIslandElement( ); const code = extractIslandCodeFromEmbed(embed); - if (!cellOutput || !code) { - return null; + if (cellOutput && code) { + return { + output: cellOutput.innerHTML, + code: code, + }; } - return { - output: cellOutput.innerHTML, - code: code, - }; + const retained = retainedIslandSources.get(embed); + return retained?.code ? retained : null; } export function createMarimoFile(app: { 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..1fc00ec5184 --- /dev/null +++ b/frontend/src/core/islands/worker/__tests__/controller.test.ts @@ -0,0 +1,173 @@ +/* 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 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(stopImplementation), { + destroy: vi.fn(), + }); + const packages = { destroy: vi.fn(), toJs: () => [] }; + 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() + .mockImplementation(() => sessions.shift()?.sessionResources), + loadPackagesFromImports, + loadedPackages: {}, + runPythonAsync: vi.fn(), + } as unknown as PyodideInterface; + + return { + ...first, + loadPackagesFromImports, + pyodide, + }; +} + +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"); + }); + + 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(); + }); + + 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/islands/worker/worker.tsx b/frontend/src/core/islands/worker/worker.tsx index 090236f9b86..f17668c6e60 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,128 @@ 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) { + (nextBridge as unknown as PyProxy).destroy(); + } else { + 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; + await self.controller.stopSession(); + activeSession = undefined; + (session?.bridge as unknown as PyProxy | undefined)?.destroy(); +} + +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 ( + 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 +179,56 @@ 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 !== 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< { diff --git a/frontend/src/core/wasm/worker/bootstrap.ts b/frontend/src/core/wasm/worker/bootstrap.ts index 8339b422258..91ae31d88eb 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,12 @@ import type { SerializedBridge, WasmController } from "./types"; import { shouldLoadDuckDBPackages } from "../utils"; const MAKE_SNAPSHOT = false; +type SessionResources = [ + bridge: PyProxy, + init: PyCallable, + packages: PyProxy, + stop: PyCallable, +]; // This class initializes the wasm environment // We would like this initialization to be parallelizable @@ -25,6 +32,9 @@ const MAKE_SNAPSHOT = false; export class DefaultWasmController implements WasmController { protected pyodide: PyodideInterface | null = null; + private packageLoadQueue = Promise.resolve(); + private sessionGeneration = 0; + private activeSessionStops = new Set(); get requirePyodide() { invariant(this.pyodide, "Pyodide not loaded"); @@ -111,6 +121,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 +137,129 @@ 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 + 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: + 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() - 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.activeSessionStops.add(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 bridgeProxy as unknown as SerializedBridge; + } - return bridge; + async stopSession(): Promise { + this.sessionGeneration += 1; + const stops = [...this.activeSessionStops]; + let hasFailure = false; + let firstFailure: unknown; + for (const stop of stops) { + try { + await stop(); + this.activeSessionStops.delete(stop); + stop.destroy(); + } catch (error) { + if (!hasFailure) { + hasFailure = true; + firstFailure = error; + } + } + } + if (hasFailure) { + throw firstFailure; + } } private async loadNotebookDeps(code: string, foundPackages: Set) {