feat: add --execute to marimo export html-wasm for session previews (#9437)

## 📝 Summary

This PR introduces `--execute` to `marimo export html-wasm` enabling
previews before running.

Example:

<img width="1158" height="995" alt="image"
src="https://github.com/user-attachments/assets/68fa947d-89c8-4314-adc9-ef350530b467"
/>

Author: dmadisetti

frontend/src/core/cells/__tests__/cells.test.ts

@@ -46,6 +46,9 @@ vi.mock("@/core/codemirror/editing/commands", () => ({
   foldAllBulk: vi.fn(),
   unfoldAllBulk: vi.fn(),
 }));
+vi.mock("@/core/wasm/utils", () => ({
+  isWasm: vi.fn(() => false),
+}));
 vi.mock("../scrollCellIntoView", async (importOriginal) => {
   const actual = await importOriginal();
   return {
@@ -3428,3 +3431,99 @@ describe("createTracebackInfoAtom", () => {
     expect(traceback).toBeUndefined();
   });
 });
+
+describe("setCells snapshot preservation", () => {
+  const CELL_A = cellId("A");
+  const CELL_B = cellId("B");
+  const newCells: CellData[] = [
+    {
+      id: CELL_A,
+      name: "a",
+      code: "1",
+      edited: false,
+      lastCodeRun: null,
+      lastExecutionTime: null,
+      config: { hide_code: false, disabled: false, column: null },
+      serializedEditorState: null,
+    },
+    {
+      id: CELL_B,
+      name: "b",
+      code: "2",
+      edited: false,
+      lastCodeRun: null,
+      lastExecutionTime: null,
+      config: { hide_code: false, disabled: false, column: null },
+      serializedEditorState: null,
+    },
+  ];
+
+  const hydratedState = () =>
+    MockNotebook.notebookState({
+      cellData: {
+        [CELL_A]: { id: CELL_A, code: "1" },
+        [CELL_B]: { id: CELL_B, code: "2" },
+      },
+      cellRuntime: {
+        [CELL_A]: {
+          output: {
+            channel: "output",
+            mimetype: "text/plain",
+            data: "hydrated-A",
+            timestamp: 0,
+          },
+        },
+        [CELL_B]: {
+          consoleOutputs: [
+            {
+              channel: "stdout",
+              mimetype: "text/plain",
+              data: "hydrated-B-stdout",
+              timestamp: 0,
+            },
+          ],
+        },
+      },
+    });
+
+  beforeEach(async () => {
+    const { isWasm } = await import("@/core/wasm/utils");
+    vi.mocked(isWasm).mockReturnValue(true);
+  });
+
+  it("preserves hydrated output in WASM", () => {
+    const next = exportedForTesting.reducer(hydratedState(), {
+      type: "setCells",
+      payload: newCells,
+    });
+
+    expect(next.cellRuntime[CELL_A].output).toMatchObject({
+      data: "hydrated-A",
+    });
+  });
+
+  it("preserves console-only hydration in WASM", () => {
+    const next = exportedForTesting.reducer(hydratedState(), {
+      type: "setCells",
+      payload: newCells,
+    });
+
+    expect(next.cellRuntime[CELL_B].consoleOutputs).toHaveLength(1);
+    expect(next.cellRuntime[CELL_B].consoleOutputs[0]).toMatchObject({
+      data: "hydrated-B-stdout",
+    });
+  });
+
+  it("resets cells with no prior runtime even in WASM", () => {
+    const empty = MockNotebook.notebookState({ cellData: {} });
+    const next = exportedForTesting.reducer(empty, {
+      type: "setCells",
+      payload: newCells,
+    });
+
+    expect(next.cellRuntime[CELL_A].output).toBeNull();
+    expect(next.cellRuntime[CELL_A].consoleOutputs).toEqual([]);
+    expect(next.cellRuntime[CELL_B].output).toBeNull();
+    expect(next.cellRuntime[CELL_B].consoleOutputs).toEqual([]);
+  });
+});

frontend/src/core/cells/cells.ts

@@ -29,6 +29,7 @@ import { isErrorMime } from "../mime";
 import type { CellConfig } from "../network/types";
 import { isRtcEnabled } from "../rtc/state";
 import { createDeepEqualAtom, store } from "../state/jotai";
+import { isWasm } from "../wasm/utils";
 import { prepareCellForExecution, transitionCell } from "./cell";
 import { documentTransactionMiddleware } from "./document-changes";
 import { CellId, SCRATCH_CELL_ID, SETUP_CELL_ID } from "./ids";
@@ -1032,8 +1033,20 @@ const {
   setCells: (state, cells: CellData[]) => {
     const cellData = Object.fromEntries(cells.map((cell) => [cell.id, cell]));
 
+    // WASM has no server-side SessionView to replay outputs, so the
+    // snapshot hydrated by notebookStateFromSession is the only source.
+    const preserveSnapshot = isWasm();
+    const runtimeFor = (cellId: CellId): CellRuntimeState => {
+      if (!preserveSnapshot) {
+        return createCellRuntimeState();
+      }
+      const prev = state.cellRuntime[cellId];
+      const hasSnapshot =
+        prev && (prev.output != null || prev.consoleOutputs.length > 0);
+      return hasSnapshot ? prev : createCellRuntimeState();
+    };
     const cellRuntime = Object.fromEntries(
-      cells.map((cell) => [cell.id, createCellRuntimeState()]),
+      cells.map((cell) => [cell.id, runtimeFor(cell.id)]),
     );
 
     return withScratchCell({

marimo/_cli/export/commands.py

@@ -32,6 +32,7 @@
     run_app_then_export_as_html,
     run_app_then_export_as_ipynb,
     run_app_then_export_as_pdf,
+    run_app_then_export_as_wasm,
 )
 from marimo._server.export._status import PDFExportStatusEvent
 from marimo._server.export.exporter import Exporter
@@ -870,11 +871,20 @@ def export_callback_impl(file_path: MarimoPath) -> ExportResult:
     default=False,
     help="Force overwrite of the output file if it already exists.",
 )
+@click.option(
+    "--execute/--no-execute",
+    default=False,
+    help=(
+        "Execute the notebook before exporting and embed outputs as a "
+        "preview. Runs in the current Python environment."
+    ),
+)
 @click.argument(
     "name",
     required=True,
     type=click.Path(exists=True, file_okay=True, dir_okay=False),
 )
+@click.argument("args", nargs=-1, type=click.UNPROCESSED)
 def html_wasm(
     name: str,
     output: Path,
@@ -884,10 +894,17 @@ def html_wasm(
     include_cloudflare: bool,
     sandbox: bool | None,
     force: bool,
+    execute: bool,
+    args: tuple[str, ...],
 ) -> None:
     """Export a notebook as a WASM-powered standalone HTML file."""
     import sys
 
+    if execute and watch:
+        raise click.UsageError(
+            "--execute and --watch cannot be used together."
+        )
+
     # Set default, if not provided
     if sandbox is None:
         from marimo._cli.sandbox import maybe_prompt_run_in_sandbox
@@ -909,8 +926,25 @@ def html_wasm(
 
     marimo_file = MarimoPath(name)
 
-    def export_callback(file_path: MarimoPath) -> ExportResult:
-        return export_as_wasm(file_path, mode, show_code=show_code)
+    if execute:
+        cli_args = parse_args(args)
+
+        def export_callback(file_path: MarimoPath) -> ExportResult:
+            return asyncio_run(
+                run_app_then_export_as_wasm(
+                    file_path,
+                    mode=mode,
+                    show_code=show_code,
+                    cli_args=cli_args,
+                    argv=list(args),
+                )
+            )
+
+        echo("Executing notebook...")
+    else:
+
+        def export_callback(file_path: MarimoPath) -> ExportResult:
+            return export_as_wasm(file_path, mode, show_code=show_code)
 
     # Export assets first
     Exporter().export_assets(out_dir)

marimo/_server/export/__init__.py

@@ -355,6 +355,60 @@ async def run_app_then_export_as_html(
     )
 
 
+async def run_app_then_export_as_wasm(
+    path: MarimoPath,
+    mode: Literal["edit", "run"],
+    show_code: bool,
+    cli_args: SerializedCLIArgs,
+    argv: list[str],
+    *,
+    asset_url: str | None = None,
+) -> ExportResult:
+    """Execute notebook and export as WASM HTML with embedded session."""
+    from marimo._session.state.serialize import (
+        serialize_notebook,
+        serialize_session_view,
+    )
+
+    file_manager = load_notebook(path.absolute_name)
+    file_manager.app.inline_layout_file()
+
+    config = get_default_config_manager(current_path=file_manager.path)
+    display_config = cast(DisplayConfig, config.get_config()["display"])
+
+    session_view, did_error = await run_app_until_completion(
+        file_manager,
+        cli_args,
+        argv=argv,
+    )
+
+    session_snapshot = serialize_session_view(
+        session_view,
+        cell_ids=file_manager.app.cell_manager.cell_ids(),
+        drop_virtual_file_outputs=True,
+    )
+    notebook_snapshot = serialize_notebook(
+        session_view, file_manager.app.cell_manager
+    )
+
+    html, filename = Exporter().export_as_wasm(
+        filename=file_manager.filename,
+        app=file_manager.app,
+        display_config=display_config,
+        code=file_manager.app.to_py(),
+        mode=mode,
+        show_code=show_code,
+        asset_url=asset_url,
+        session_snapshot=session_snapshot,
+        notebook_snapshot=notebook_snapshot,
+    )
+    return ExportResult(
+        contents=html,
+        download_filename=filename,
+        did_error=did_error,
+    )
+
+
 async def export_as_html_without_execution(
     path: MarimoPath,
     include_code: bool,

marimo/_server/export/exporter.py

@@ -328,6 +328,8 @@ def export_as_wasm(
         mode: Literal["edit", "run"],
         show_code: bool,
         asset_url: str | None = None,
+        session_snapshot: NotebookSessionV1 | None = None,
+        notebook_snapshot: NotebookV1 | None = None,
     ) -> tuple[str, str]:
         """Export notebook as a WASM-powered standalone HTML file."""
         index_html = get_html_contents()
@@ -350,6 +352,8 @@ def export_as_wasm(
             code=code,
             asset_url=asset_url,
             show_code=show_code,
+            session_snapshot=session_snapshot,
+            notebook_snapshot=notebook_snapshot,
         )
 
         download_filename = get_download_filename(filename, "wasm.html")

marimo/_server/templates/templates.py

@@ -503,6 +503,8 @@ def wasm_notebook_template(
     code: str,
     show_code: bool,
     asset_url: str | None = None,
+    session_snapshot: NotebookSessionV1 | None = None,
+    notebook_snapshot: NotebookV1 | None = None,
 ) -> str:
     """Template for WASM notebooks."""
     import re
@@ -536,6 +538,8 @@ def wasm_notebook_template(
             version=version,
             show_app_code=show_code,
             runtime_config=None,
+            session_snapshot=session_snapshot,
+            notebook_snapshot=notebook_snapshot,
         ),
     )