diff --git a/frontend/src/mount.tsx b/frontend/src/mount.tsx index e836e07b2d6..ed495f278b3 100644 --- a/frontend/src/mount.tsx +++ b/frontend/src/mount.tsx @@ -20,12 +20,14 @@ import { import { Logger } from "@/utils/Logger"; import { ErrorBoundary } from "./components/editor/boundary/ErrorBoundary"; import { notebookAtom } from "./core/cells/cells"; +import type { UIElementId } from "./core/cells/ids"; import { notebookStateFromSession } from "./core/cells/session"; import { parseAppConfig, parseConfigOverrides, parseUserConfig, } from "./core/config/config-schema"; +import { UI_ELEMENT_REGISTRY } from "./core/dom/uiregistry"; import { MarimoApp, preloadPage } from "./core/MarimoApp"; import { type AppMode, initialModeAtom, viewStateAtom } from "./core/mode"; import { cleanupAuthQueryParams } from "./core/network/auth"; @@ -341,6 +343,16 @@ function initStore(options: unknown) { if (notebook) { store.set(notebookAtom, notebook); } + + // Restore widget state from session snapshot in static mode + if (isStaticNotebook()) { + const uiValues = parsedOptions.data.session?.ui_values; + if (uiValues) { + for (const [objectId, value] of Object.entries(uiValues)) { + UI_ELEMENT_REGISTRY.set(objectId as UIElementId, value); + } + } + } } /** diff --git a/marimo/_schemas/generated/session.yaml b/marimo/_schemas/generated/session.yaml index 82fd56459de..5149b2a131a 100644 --- a/marimo/_schemas/generated/session.yaml +++ b/marimo/_schemas/generated/session.yaml @@ -79,6 +79,10 @@ components: type: array metadata: $ref: '#/components/schemas/NotebookSessionMetadata' + ui_values: + additionalProperties: {} + nullable: true + type: object version: type: string required: diff --git a/marimo/_schemas/session.py b/marimo/_schemas/session.py index 6c664e2add0..707a375c710 100644 --- a/marimo/_schemas/session.py +++ b/marimo/_schemas/session.py @@ -98,6 +98,8 @@ class NotebookSessionV1(BaseDict): metadata: NotebookSessionMetadata # The cells in the notebook cells: list[Cell] + # Widget/UI element values at time of export + ui_values: Optional[dict[str, Any]] # In future, we may want to add # - variables diff --git a/marimo/_session/state/serialize.py b/marimo/_session/state/serialize.py index 9f664df45f0..757d33020b1 100644 --- a/marimo/_session/state/serialize.py +++ b/marimo/_session/state/serialize.py @@ -163,6 +163,7 @@ def serialize_session_view( script_metadata_hash=script_metadata_hash, ), cells=cells, + ui_values=dict(view.ui_values), ) diff --git a/packages/openapi/src/session.ts b/packages/openapi/src/session.ts index d08655759d4..1c9b8dcbfa7 100644 --- a/packages/openapi/src/session.ts +++ b/packages/openapi/src/session.ts @@ -40,6 +40,9 @@ export interface components { NotebookSessionV1: { cells: components["schemas"]["Cell"][]; metadata: components["schemas"]["NotebookSessionMetadata"]; + ui_values?: { + [key: string]: unknown; + } | null; version: string; }; StreamMediaOutput: { diff --git a/tests/_session/state/snapshots/basic_session.json b/tests/_session/state/snapshots/basic_session.json index 19d57390340..72652c69c37 100644 --- a/tests/_session/state/snapshots/basic_session.json +++ b/tests/_session/state/snapshots/basic_session.json @@ -18,5 +18,6 @@ ], "console": [] } - ] -} + ], + "ui_values": {} +} \ No newline at end of file diff --git a/tests/_session/state/snapshots/console_session.json b/tests/_session/state/snapshots/console_session.json index 0e54cfa1793..945ef55c8ba 100644 --- a/tests/_session/state/snapshots/console_session.json +++ b/tests/_session/state/snapshots/console_session.json @@ -24,5 +24,6 @@ } ] } - ] + ], + "ui_values": {} } \ No newline at end of file diff --git a/tests/_session/state/snapshots/error_session.json b/tests/_session/state/snapshots/error_session.json index 2c05bfe7074..8c3b6144448 100644 --- a/tests/_session/state/snapshots/error_session.json +++ b/tests/_session/state/snapshots/error_session.json @@ -18,5 +18,6 @@ ], "console": [] } - ] -} + ], + "ui_values": {} +} \ No newline at end of file diff --git a/tests/_session/state/snapshots/mime_bundle_session.json b/tests/_session/state/snapshots/mime_bundle_session.json index 7420a993ab2..287f28f072a 100644 --- a/tests/_session/state/snapshots/mime_bundle_session.json +++ b/tests/_session/state/snapshots/mime_bundle_session.json @@ -21,5 +21,6 @@ ], "console": [] } - ] -} + ], + "ui_values": {} +} \ No newline at end of file diff --git a/tests/_session/state/snapshots/mixed_error_session.json b/tests/_session/state/snapshots/mixed_error_session.json index 711dcc6edd2..610a67c2866 100644 --- a/tests/_session/state/snapshots/mixed_error_session.json +++ b/tests/_session/state/snapshots/mixed_error_session.json @@ -24,5 +24,6 @@ ], "console": [] } - ] + ], + "ui_values": {} } \ No newline at end of file