@marimo-team/islands 0.25.3-dev2 → 0.25.3-dev200

package/src/core/cells/__tests__/session.test.ts

@@ -6,11 +6,10 @@ import { beforeEach, describe, expect, it, vi } from "vitest";
 import { Mocks } from "@/__mocks__/common";
 import { cellId } from "@/__tests__/branded";
 import { parseOutline } from "@/core/dom/outline";
-import { MultiColumn, visibleForTesting } from "@/utils/id-tree";
 import { invariant } from "@/utils/invariant";
 import { Logger } from "@/utils/Logger";
 import { SETUP_CELL_ID } from "../ids";
-import { notebookStateFromSession } from "../session";
+import { readExport } from "../session";
 
 // Mock dependencies
 vi.mock("@/utils/Logger", () => ({ Logger: Mocks.quietLogger() }));
@@ -25,10 +24,16 @@ type NotebookCell = api.Notebook["NotebookV1"]["cells"][0];
 // Test constants
 const CELL_1 = cellId("cell-1");
 
-describe("notebookStateFromSession", () => {
+/** The snapshot's record of a cell. */
+function cellIn(result: ReturnType<typeof readExport>, id: string) {
+  const cell = result.snapshot?.cells.find((cell) => cell.id === id);
+  invariant(cell, `no cell ${id}`);
+  return cell;
+}
+
+describe("readExport", () => {
   beforeEach(() => {
     vi.clearAllMocks();
-    visibleForTesting.reset();
   });
 
   // Test data factories
@@ -46,10 +51,17 @@ describe("notebookStateFromSession", () => {
 
   const createNotebookCell = (
     id: string,
-    code: string | null = null,
-    name: string | null = null,
-    config: NotebookCell["config"] | null = null,
-    code_hash: string | null = null,
+    {
+      code = null,
+      name = null,
+      config = null,
+      code_hash = null,
+    }: {
+      code?: string | null;
+      name?: string | null;
+      config?: NotebookCell["config"] | null;
+      code_hash?: string | null;
+    } = {},
   ): NotebookCell => ({
     id,
     code,
@@ -81,72 +93,113 @@ describe("notebookStateFromSession", () => {
 
   describe("validation", () => {
     it("logs error for both session and notebook null/undefined", () => {
-      const result = notebookStateFromSession(null, null);
-      expect(result).toBeNull();
+      const result = readExport({
+        session: null,
+        notebook: null,
+        filename: null,
+        appConfig: {},
+      });
+      expect(result.snapshot).toBeNull();
     });
 
     it("logs error when session and notebook have different cells", () => {
       const session = createSession([createSessionCell("cell-1")]);
-      const notebook = createNotebook([createNotebookCell("cell-2")]);
+      const notebook = createNotebook([createNotebookCell("cell-2", {})]);
 
-      const result = notebookStateFromSession(session, notebook);
+      const result = readExport({
+        session: session,
+        notebook: notebook,
+        filename: null,
+        appConfig: {},
+      });
 
       expect(Logger.warn).toHaveBeenCalledWith(
         "Session and notebook have different cells, attempted merge.",
       );
       // Should have the same cell structure as notebook
-      expect(result).not.toBeNull();
-      invariant(result, "result is null");
-      expect(result.cellIds.inOrderIds).toEqual(
-        MultiColumn.from([["cell-2"]]).inOrderIds,
-      );
+      invariant(result.snapshot, "no snapshot");
+      expect(result.snapshot.document.order).toEqual(["cell-2"]);
     });
   });
 
   describe("null/undefined inputs", () => {
     it("returns null when both session and notebook are null", () => {
-      const result = notebookStateFromSession(null, null);
-      expect(result).toBeNull();
+      const result = readExport({
+        session: null,
+        notebook: null,
+        filename: null,
+        appConfig: {},
+      });
+      expect(result.snapshot).toBeNull();
     });
 
     it("returns null when both session and notebook are undefined", () => {
-      const result = notebookStateFromSession(undefined, undefined);
-      expect(result).toBeNull();
+      const result = readExport({
+        session: undefined,
+        notebook: undefined,
+        filename: null,
+        appConfig: {},
+      });
+      expect(result.snapshot).toBeNull();
     });
 
     it("returns null when session is null and notebook is undefined", () => {
-      const result = notebookStateFromSession(null, undefined);
-      expect(result).toBeNull();
+      const result = readExport({
+        session: null,
+        notebook: undefined,
+        filename: null,
+        appConfig: {},
+      });
+      expect(result.snapshot).toBeNull();
     });
   });
 
   describe("session only scenarios", () => {
     it("creates state from session with single cell", () => {
       const session = createSession([createSessionCell("cell-1")]);
-      const result = notebookStateFromSession(session, null);
+      const result = readExport({
+        session: session,
+        notebook: null,
+        filename: null,
+        appConfig: {},
+      });
 
-      expect(result).not.toBeNull();
-      invariant(result, "result is null");
-      expect(result.cellIds.inOrderIds).toEqual(
-        MultiColumn.from([[CELL_1]]).inOrderIds,
-      );
-      expect(result.cellData[CELL_1].code).toBe("");
+      invariant(result.snapshot, "no snapshot");
+      expect(result.snapshot.document.order).toEqual([CELL_1]);
+      // The session knows outputs, not source; the host withholds it as null.
+      expect(cellIn(result, CELL_1).code).toBeNull();
       expect(result.cellRuntime[CELL_1]).toBeDefined();
     });
 
+    it("carries the page's filename and app config into the snapshot", () => {
+      const session = createSession([createSessionCell("cell-1")]);
+      const result = readExport({
+        session,
+        notebook: null,
+        filename: "notebook.py",
+        appConfig: { width: "full" },
+      });
+
+      invariant(result.snapshot, "no snapshot");
+      expect(result.snapshot.notebook.name).toBe("notebook.py");
+      expect(result.snapshot.document.app_config).toEqual({ width: "full" });
+    });
+
     it("creates state from session with multiple cells", () => {
       const session = createSession([
         createSessionCell("cell-1"),
         createSessionCell("cell-2"),
       ]);
-      const result = notebookStateFromSession(session, null);
+      const result = readExport({
+        session: session,
+        notebook: null,
+        filename: null,
+        appConfig: {},
+      });
 
-      expect(result).not.toBeNull();
-      invariant(result, "result is null");
-      expect(result.cellIds.inOrderIds).toEqual(
-        MultiColumn.from([["cell-1", "cell-2"]]).inOrderIds,
-      );
-      expect(Object.keys((result as any).cellData)).toEqual([
+      invariant(result.snapshot, "no snapshot");
+      expect(result.snapshot.document.order).toEqual(["cell-1", "cell-2"]);
+      expect(result.snapshot.cells.map((cell) => cell.id)).toEqual([
         "cell-1",
         "cell-2",
       ]);
@@ -166,16 +219,23 @@ describe("notebookStateFromSession", () => {
       const session = createSession([
         createSessionCell("cell-1", [errorOutput]),
       ]);
-      const result = notebookStateFromSession(session, null);
+      const result = readExport({
+        session: session,
+        notebook: null,
+        filename: null,
+        appConfig: {},
+      });
 
-      expect(result).not.toBeNull();
-      invariant(result, "result is null");
+      invariant(result.snapshot, "no snapshot");
       expect(result.cellRuntime[CELL_1].output).toEqual({
         channel: "marimo-error",
         data: [
           {
-            type: "unknown",
+            type: "exception",
+            exception_type: "ValueError",
             msg: "Something went wrong",
+            raising_cell: null,
+            traceback: "Traceback line 1",
           },
         ],
         mimetype: "application/vnd.marimo+error",
@@ -183,6 +243,57 @@ describe("notebookStateFromSession", () => {
       });
     });
 
+    it("restores matching code from the session", () => {
+      const code = "raise ValueError('cached failure')";
+      const session = createSession([
+        createSessionCell(
+          "cell-1",
+          [
+            {
+              type: "error",
+              ename: "ValueError",
+              evalue: "cached failure",
+              traceback: ["Traceback\n", "ValueError: cached failure"],
+              code,
+            },
+          ],
+          [
+            {
+              type: "stream",
+              name: "stderr",
+              text: "ValueError: cached <failure>",
+              mimetype: "application/vnd.marimo+traceback",
+              traceback_format: "text",
+              code,
+            },
+          ],
+        ),
+      ]);
+      const state = readExport({
+        session: session,
+        notebook: null,
+        filename: null,
+        appConfig: {},
+      });
+      invariant(state.snapshot, "no snapshot");
+      expect(state.cellRuntime[CELL_1].output).toMatchObject({
+        code,
+        data: [
+          {
+            type: "exception",
+            exception_type: "ValueError",
+            msg: "cached failure",
+            traceback: "Traceback\nValueError: cached failure",
+          },
+        ],
+      });
+      expect(state.cellRuntime[CELL_1].consoleOutputs[0].code).toBe(code);
+      expect(state.cellRuntime[CELL_1].consoleOutputs[0]).toMatchObject({
+        traceback_format: "text",
+        data: "ValueError: cached <failure>",
+      });
+    });
+
     it("handles data output in session cell", () => {
       const dataOutput = {
         type: "data" as const,
@@ -193,10 +304,14 @@ describe("notebookStateFromSession", () => {
       const session = createSession([
         createSessionCell("cell-1", [dataOutput]),
       ]);
-      const result = notebookStateFromSession(session, null);
+      const result = readExport({
+        session: session,
+        notebook: null,
+        filename: null,
+        appConfig: {},
+      });
 
-      expect(result).not.toBeNull();
-      invariant(result, "result is null");
+      invariant(result.snapshot, "no snapshot");
       expect(result.cellRuntime[CELL_1].output).toEqual({
         channel: "output",
         data: "Hello World",
@@ -223,10 +338,14 @@ describe("notebookStateFromSession", () => {
       const session = createSession([
         createSessionCell("cell-1", [], consoleOutputs),
       ]);
-      const result = notebookStateFromSession(session, null);
+      const result = readExport({
+        session: session,
+        notebook: null,
+        filename: null,
+        appConfig: {},
+      });
 
-      expect(result).not.toBeNull();
-      invariant(result, "result is null");
+      invariant(result.snapshot, "no snapshot");
       expect(result.cellRuntime[CELL_1].consoleOutputs).toEqual([
         {
           channel: "stdout",
@@ -254,7 +373,12 @@ describe("notebookStateFromSession", () => {
       const session = createSession([
         createSessionCell("cell-1", [dataOutput]),
       ]);
-      const result = notebookStateFromSession(session, null);
+      const result = readExport({
+        session: session,
+        notebook: null,
+        filename: null,
+        appConfig: {},
+      });
 
       expect(parseOutline).toHaveBeenCalledWith({
         channel: "output",
@@ -262,16 +386,21 @@ describe("notebookStateFromSession", () => {
         mimetype: "text/html",
         timestamp: 0,
       });
-      invariant(result, "result is null");
+      invariant(result.snapshot, "no snapshot");
       expect(result.cellRuntime[CELL_1].outline).toBe(mockOutline);
     });
 
     it("sets outline to null when no output exists", () => {
       const session = createSession([createSessionCell("cell-1")]);
-      const result = notebookStateFromSession(session, null);
+      const result = readExport({
+        session: session,
+        notebook: null,
+        filename: null,
+        appConfig: {},
+      });
 
       expect(parseOutline).not.toHaveBeenCalled();
-      invariant(result, "result is null");
+      invariant(result.snapshot, "no snapshot");
       expect(result.cellRuntime[CELL_1].outline).toBeNull();
     });
   });
@@ -279,46 +408,59 @@ describe("notebookStateFromSession", () => {
   describe("notebook only scenarios", () => {
     it("creates state from notebook with single cell", () => {
       const notebook = createNotebook([
-        createNotebookCell("cell-1", "print('hello')", "my_cell"),
+        createNotebookCell("cell-1", {
+          code: "print('hello')",
+          name: "my_cell",
+        }),
       ]);
-      const result = notebookStateFromSession(null, notebook);
+      const result = readExport({
+        session: null,
+        notebook: notebook,
+        filename: null,
+        appConfig: {},
+      });
 
-      expect(result).not.toBeNull();
-      invariant(result, "result is null");
-      expect(result.cellIds.inOrderIds).toEqual(
-        MultiColumn.from([[CELL_1]]).inOrderIds,
-      );
-      expect(result.cellData[CELL_1]).toEqual({
-        id: "cell-1",
+      invariant(result.snapshot, "no snapshot");
+      expect(result.snapshot.document.order).toEqual([CELL_1]);
+      expect(cellIn(result, CELL_1)).toEqual({
+        id: CELL_1,
+        notebook_id: "embedded",
         name: "my_cell",
         code: "print('hello')",
-        edited: false,
-        lastCodeRun: null,
-        lastExecutionTime: null,
         config: {
           hide_code: false,
           disabled: false,
           column: null,
           expand_output: false,
         },
-        serializedEditorState: null,
+        etag: `embedded:${CELL_1}`,
+        code_hash: `embedded:${CELL_1}:code`,
+        version: 1,
       });
     });
 
     it("creates state from notebook with custom config", () => {
       const notebook = createNotebook([
-        createNotebookCell("cell-1", "code", "name", {
-          hide_code: true,
-          disabled: true,
-          column: 1,
-          expand_output: true,
+        createNotebookCell("cell-1", {
+          code: "code",
+          name: "name",
+          config: {
+            hide_code: true,
+            disabled: true,
+            column: 1,
+            expand_output: true,
+          },
         }),
       ]);
-      const result = notebookStateFromSession(null, notebook);
+      const result = readExport({
+        session: null,
+        notebook: notebook,
+        filename: null,
+        appConfig: {},
+      });
 
-      expect(result).not.toBeNull();
-      invariant(result, "result is null");
-      expect(result.cellData[CELL_1].config).toEqual({
+      invariant(result.snapshot, "no snapshot");
+      expect(cellIn(result, CELL_1).config).toEqual({
         hide_code: true,
         disabled: true,
         column: 1,
@@ -339,24 +481,28 @@ describe("notebookStateFromSession", () => {
         },
       };
       const notebook = createNotebook([notebookCell]);
-      const result = notebookStateFromSession(null, notebook);
+      const result = readExport({
+        session: null,
+        notebook: notebook,
+        filename: null,
+        appConfig: {},
+      });
 
-      expect(result).not.toBeNull();
-      invariant(result, "result is null");
-      expect(result.cellData[CELL_1]).toEqual({
-        id: "cell-1",
-        name: "",
-        code: "",
-        edited: false,
-        lastCodeRun: null,
-        lastExecutionTime: null,
+      invariant(result.snapshot, "no snapshot");
+      expect(cellIn(result, CELL_1)).toEqual({
+        id: CELL_1,
+        notebook_id: "embedded",
+        name: "_",
+        code: null,
         config: {
           hide_code: false,
           disabled: false,
           column: null,
           expand_output: false,
         },
-        serializedEditorState: null,
+        etag: `embedded:${CELL_1}`,
+        code_hash: `embedded:${CELL_1}:code`,
+        version: 1,
       });
     });
 
@@ -373,27 +519,29 @@ describe("notebookStateFromSession", () => {
         },
       };
       const notebook = createNotebook([notebookCell as any]);
-      const result = notebookStateFromSession(null, notebook);
+      const result = readExport({
+        session: null,
+        notebook: notebook,
+        filename: null,
+        appConfig: {},
+      });
 
-      expect(result).not.toBeNull();
-      invariant(result, "result is null");
-      expect(result.cellIds.inOrderIds).toEqual(
-        MultiColumn.from([[SETUP_CELL_ID]]).inOrderIds,
-      );
-      expect(result.cellData[SETUP_CELL_ID]).toEqual({
+      invariant(result.snapshot, "no snapshot");
+      expect(result.snapshot.document.order).toEqual([SETUP_CELL_ID]);
+      expect(cellIn(result, SETUP_CELL_ID)).toEqual({
         id: SETUP_CELL_ID,
+        notebook_id: "embedded",
         name: "setup",
         code: "import marimo as mo",
-        edited: false,
-        lastCodeRun: null,
-        lastExecutionTime: null,
         config: {
           hide_code: false,
           disabled: false,
           column: null,
           expand_output: false,
         },
-        serializedEditorState: null,
+        etag: `embedded:${SETUP_CELL_ID}`,
+        code_hash: `embedded:${SETUP_CELL_ID}:code`,
+        version: 1,
       });
     });
   });
@@ -415,26 +563,33 @@ describe("notebookStateFromSession", () => {
         ),
       ]);
       const notebook = createNotebook([
-        createNotebookCell("cell-1", "print('hello')", "my_cell"),
+        createNotebookCell("cell-1", {
+          code: "print('hello')",
+          name: "my_cell",
+        }),
       ]);
-      const result = notebookStateFromSession(session, notebook);
+      const result = readExport({
+        session: session,
+        notebook: notebook,
+        filename: null,
+        appConfig: {},
+      });
 
-      expect(result).not.toBeNull();
-      invariant(result, "result is null");
-      expect(result.cellData[CELL_1]).toEqual({
-        id: "cell-1",
+      invariant(result.snapshot, "no snapshot");
+      expect(cellIn(result, CELL_1)).toEqual({
+        id: CELL_1,
+        notebook_id: "embedded",
         name: "my_cell",
         code: "print('hello')",
-        edited: false,
-        lastCodeRun: null,
-        lastExecutionTime: null,
         config: {
           hide_code: false,
           disabled: false,
           column: null,
           expand_output: false,
         },
-        serializedEditorState: null,
+        etag: `embedded:${CELL_1}`,
+        code_hash: `embedded:${CELL_1}:code`,
+        version: 1,
       });
       // Null result means no output.
       expect(result.cellRuntime[CELL_1].consoleOutputs).toEqual([]);
@@ -446,17 +601,19 @@ describe("notebookStateFromSession", () => {
         createSessionCell("cell-2"),
       ]);
       const notebook = createNotebook([
-        createNotebookCell("cell-2"),
-        createNotebookCell("cell-1"),
+        createNotebookCell("cell-2", {}),
+        createNotebookCell("cell-1", {}),
       ]);
 
-      const result = notebookStateFromSession(session, notebook);
+      const result = readExport({
+        session: session,
+        notebook: notebook,
+        filename: null,
+        appConfig: {},
+      });
 
-      expect(result).not.toBeNull();
-      invariant(result, "result is null");
-      expect(result.cellIds.inOrderIds).toEqual(
-        MultiColumn.from([["cell-2", "cell-1"]]).inOrderIds,
-      );
+      invariant(result.snapshot, "no snapshot");
+      expect(result.snapshot.document.order).toEqual(["cell-2", "cell-1"]);
     });
 
     it("creates state when session and notebook have identical cell arrays", () => {
@@ -465,52 +622,73 @@ describe("notebookStateFromSession", () => {
         createSessionCell("cell-2"),
       ]);
       const notebook = createNotebook([
-        createNotebookCell("cell-1"),
-        createNotebookCell("cell-2"),
+        createNotebookCell("cell-1", {}),
+        createNotebookCell("cell-2", {}),
       ]);
 
-      const result = notebookStateFromSession(session, notebook);
+      const result = readExport({
+        session: session,
+        notebook: notebook,
+        filename: null,
+        appConfig: {},
+      });
 
       expect(Logger.error).not.toHaveBeenCalled();
-      expect(result).not.toBeNull();
-      invariant(result, "result is null");
-      expect(result.cellIds.inOrderIds).toEqual(
-        MultiColumn.from([["cell-1", "cell-2"]]).inOrderIds,
-      );
+      invariant(result.snapshot, "no snapshot");
+      expect(result.snapshot.document.order).toEqual(["cell-1", "cell-2"]);
     });
   });
 
   describe("empty cells scenarios", () => {
     it("warns and returns null when session has empty cells array", () => {
       const session = createSession([]);
-      const result = notebookStateFromSession(session, null);
+      const result = readExport({
+        session: session,
+        notebook: null,
+        filename: null,
+        appConfig: {},
+      });
 
-      expect(result).toBeNull();
+      expect(result.snapshot).toBeNull();
     });
 
     it("warns and returns null when notebook has empty cells array", () => {
       const notebook = createNotebook([]);
-      const result = notebookStateFromSession(null, notebook);
+      const result = readExport({
+        session: null,
+        notebook: notebook,
+        filename: null,
+        appConfig: {},
+      });
 
-      expect(result).toBeNull();
+      expect(result.snapshot).toBeNull();
     });
 
     it("warns and returns null when both have empty cells arrays", () => {
       const session = createSession([]);
       const notebook = createNotebook([]);
-      const result = notebookStateFromSession(session, notebook);
+      const result = readExport({
+        session: session,
+        notebook: notebook,
+        filename: null,
+        appConfig: {},
+      });
 
-      expect(result).toBeNull();
+      expect(result.snapshot).toBeNull();
     });
   });
 
   describe("edge cases", () => {
     it("handles session cell with no outputs", () => {
       const session = createSession([createSessionCell("cell-1", [])]);
-      const result = notebookStateFromSession(session, null);
+      const result = readExport({
+        session: session,
+        notebook: null,
+        filename: null,
+        appConfig: {},
+      });
 
-      expect(result).not.toBeNull();
-      invariant(result, "result is null");
+      invariant(result.snapshot, "no snapshot");
       expect(result.cellRuntime[CELL_1].output).toBeNull();
       expect(result.cellRuntime[CELL_1].consoleOutputs).toEqual([]);
       expect(result.cellRuntime[CELL_1].outline).toBeNull();
@@ -518,10 +696,14 @@ describe("notebookStateFromSession", () => {
 
     it("handles session cell with empty console array", () => {
       const session = createSession([createSessionCell("cell-1", [], [])]);
-      const result = notebookStateFromSession(session, null);
+      const result = readExport({
+        session: session,
+        notebook: null,
+        filename: null,
+        appConfig: {},
+      });
 
-      expect(result).not.toBeNull();
-      invariant(result, "result is null");
+      invariant(result.snapshot, "no snapshot");
       expect(result.cellRuntime[CELL_1].consoleOutputs).toEqual([]);
     });
 
@@ -531,10 +713,14 @@ describe("notebookStateFromSession", () => {
         { type: "data" as const, data: { "text/plain": "Second" } },
       ];
       const session = createSession([createSessionCell("cell-1", outputs)]);
-      const result = notebookStateFromSession(session, null);
+      const result = readExport({
+        session: session,
+        notebook: null,
+        filename: null,
+        appConfig: {},
+      });
 
-      expect(result).not.toBeNull();
-      invariant(result, "result is null");
+      invariant(result.snapshot, "no snapshot");
       expect(result.cellRuntime[CELL_1].output).toEqual({
         channel: "output",
         data: "First",
@@ -554,10 +740,14 @@ describe("notebookStateFromSession", () => {
       const session = createSession([
         createSessionCell("cell-1", [dataOutput]),
       ]);
-      const result = notebookStateFromSession(session, null);
+      const result = readExport({
+        session: session,
+        notebook: null,
+        filename: null,
+        appConfig: {},
+      });
 
-      expect(result).not.toBeNull();
-      invariant(result, "result is null");
+      invariant(result.snapshot, "no snapshot");
       expect(result.cellRuntime[CELL_1].output).toEqual({
         channel: "output",
         data: "Plain text",
@@ -574,10 +764,14 @@ describe("notebookStateFromSession", () => {
       const session = createSession([
         createSessionCell("cell-1", [dataOutput]),
       ]);
-      const result = notebookStateFromSession(session, null);
+      const result = readExport({
+        session: session,
+        notebook: null,
+        filename: null,
+        appConfig: {},
+      });
 
-      expect(result).not.toBeNull();
-      invariant(result, "result is null");
+      invariant(result.snapshot, "no snapshot");
       expect(result.cellRuntime[CELL_1].output).toEqual({
         channel: "output",
         data: undefined,
@@ -588,22 +782,29 @@ describe("notebookStateFromSession", () => {
 
     it("returns correct structure with all expected properties", () => {
       const session = createSession([createSessionCell("cell-1")]);
-      const result = notebookStateFromSession(session, null);
-
-      expect(result).not.toBeNull();
-      invariant(result, "result is null");
-      expect(result).toHaveProperty("cellIds");
-      expect(result).toHaveProperty("cellData");
-      expect(result).toHaveProperty("cellRuntime");
-      expect(result).toHaveProperty("cellHandles");
-      expect(result).toHaveProperty("history");
-      expect(result).toHaveProperty("scrollKey");
-      expect(result).toHaveProperty("cellLogs");
-
-      expect(result.cellHandles).toEqual({});
-      expect(result.history).toEqual([]);
-      expect(result.scrollKey).toBeNull();
-      expect(result.cellLogs).toEqual([]);
+      const result = readExport({
+        session: session,
+        notebook: null,
+        filename: null,
+        appConfig: {},
+      });
+
+      invariant(result.snapshot, "no snapshot");
+      expect(result.snapshot.notebook).toEqual({
+        id: "embedded",
+        project_id: "embedded",
+        path: null,
+        name: "",
+        updated_at: null,
+      });
+      expect(result.snapshot.document).toMatchObject({
+        notebook_id: "embedded",
+        order: result.snapshot.cells.map((cell) => cell.id),
+        version: 1,
+      });
+      expect(Object.keys(result.cellRuntime)).toEqual(
+        result.snapshot.document.order,
+      );
     });
   });
 
@@ -620,12 +821,42 @@ describe("notebookStateFromSession", () => {
       };
       // Create notebook cells with code values 'a' through 'f'
       const notebookCells = [
-        createNotebookCell("cell-a", "a", null, null, hashes.a),
-        createNotebookCell("cell-b", "b", null, null, hashes.b),
-        createNotebookCell("cell-c", "c", null, null, hashes.c),
-        createNotebookCell("cell-d", "d", null, null, hashes.d),
-        createNotebookCell("cell-e", "e", null, null, hashes.e),
-        createNotebookCell("cell-f", "f", null, null, hashes.f),
+        createNotebookCell("cell-a", {
+          code: "a",
+          name: null,
+          config: null,
+          code_hash: hashes.a,
+        }),
+        createNotebookCell("cell-b", {
+          code: "b",
+          name: null,
+          config: null,
+          code_hash: hashes.b,
+        }),
+        createNotebookCell("cell-c", {
+          code: "c",
+          name: null,
+          config: null,
+          code_hash: hashes.c,
+        }),
+        createNotebookCell("cell-d", {
+          code: "d",
+          name: null,
+          config: null,
+          code_hash: hashes.d,
+        }),
+        createNotebookCell("cell-e", {
+          code: "e",
+          name: null,
+          config: null,
+          code_hash: hashes.e,
+        }),
+        createNotebookCell("cell-f", {
+          code: "f",
+          name: null,
+          config: null,
+          code_hash: hashes.f,
+        }),
       ];
       // Create session cells with matching code_hashes
       const sessionCells = [
@@ -668,19 +899,26 @@ describe("notebookStateFromSession", () => {
       ];
       const session = createSession(sessionCells);
       const notebook = createNotebook(notebookCells);
-      const result = notebookStateFromSession(session, notebook);
-      expect(result).not.toBeNull();
-      invariant(result, "result is null");
+      const result = readExport({
+        session: session,
+        notebook: notebook,
+        filename: null,
+        appConfig: {},
+      });
+      invariant(result.snapshot, "no snapshot");
       // Should have all cell IDs in order
-      expect(result.cellIds.inOrderIds).toEqual(
-        MultiColumn.from([
-          ["cell-a", "cell-b", "cell-c", "cell-d", "cell-e", "cell-f"],
-        ]).inOrderIds,
-      );
+      expect(result.snapshot.document.order).toEqual([
+        "cell-a",
+        "cell-b",
+        "cell-c",
+        "cell-d",
+        "cell-e",
+        "cell-f",
+      ]);
       // Should have correct code and output for each cell
       for (const code of ["a", "b", "c", "d", "e", "f"]) {
         const cid = cellId(`cell-${code}`);
-        expect(result.cellData[cid].code).toBe(code);
+        expect(cellIn(result, cid).code).toBe(code);
         expect(result.cellRuntime[cid].output).toEqual({
           channel: "output",
           data: `${code.toUpperCase()}!`,
@@ -738,32 +976,64 @@ describe("notebookStateFromSession", () => {
 
       // Notebook has cells with code 'a', 'c', 'z', 'e', 'g'
       const notebookCells = [
-        createNotebookCell("cell-a", "a", null, null, hashes.a),
-        createNotebookCell("cell-c", "c", null, null, hashes.c),
-        createNotebookCell("cell-z", "z", null, null, hashes.z),
-        createNotebookCell("cell-e", "e", null, null, hashes.e),
-        createNotebookCell("cell-g", "g", null, null, hashes.g),
+        createNotebookCell("cell-a", {
+          code: "a",
+          name: null,
+          config: null,
+          code_hash: hashes.a,
+        }),
+        createNotebookCell("cell-c", {
+          code: "c",
+          name: null,
+          config: null,
+          code_hash: hashes.c,
+        }),
+        createNotebookCell("cell-z", {
+          code: "z",
+          name: null,
+          config: null,
+          code_hash: hashes.z,
+        }),
+        createNotebookCell("cell-e", {
+          code: "e",
+          name: null,
+          config: null,
+          code_hash: hashes.e,
+        }),
+        createNotebookCell("cell-g", {
+          code: "g",
+          name: null,
+          config: null,
+          code_hash: hashes.g,
+        }),
       ];
 
       const session = createSession(sessionCells);
       const notebook = createNotebook(notebookCells);
-      const result = notebookStateFromSession(session, notebook);
+      const result = readExport({
+        session: session,
+        notebook: notebook,
+        filename: null,
+        appConfig: {},
+      });
 
-      expect(result).not.toBeNull();
-      invariant(result, "result is null");
+      invariant(result.snapshot, "no snapshot");
 
       // Should have notebook cell IDs in order (notebook is canonical)
-      expect(result.cellIds.inOrderIds).toEqual(
-        MultiColumn.from([["cell-a", "cell-c", "cell-z", "cell-e", "cell-g"]])
-          .inOrderIds,
-      );
+      expect(result.snapshot.document.order).toEqual([
+        "cell-a",
+        "cell-c",
+        "cell-z",
+        "cell-e",
+        "cell-g",
+      ]);
 
       // Should have correct code for each cell
-      expect(result.cellData[cellId("cell-a")].code).toBe("a");
-      expect(result.cellData[cellId("cell-c")].code).toBe("c");
-      expect(result.cellData[cellId("cell-z")].code).toBe("z");
-      expect(result.cellData[cellId("cell-e")].code).toBe("e");
-      expect(result.cellData[cellId("cell-g")].code).toBe("g");
+      expect(cellIn(result, "cell-a").code).toBe("a");
+      expect(cellIn(result, "cell-c").code).toBe("c");
+      expect(cellIn(result, "cell-z").code).toBe("z");
+      expect(cellIn(result, "cell-e").code).toBe("e");
+      expect(cellIn(result, "cell-g").code).toBe("g");
 
       // Should have session outputs for matching cells (a, c, e)
       expect(result.cellRuntime[cellId("cell-a")].output).toEqual({
@@ -859,35 +1129,48 @@ describe("notebookStateFromSession", () => {
 
       // Notebook has cells: cell-1 (1 / 0, no hash), cell-2 (mo.md("Hello"), hash=moMd), cell-3 (slider, hash=slider)
       const notebookCells = [
-        createNotebookCell("cell-1", "1 / 0", null, null, null), // no hash
-        createNotebookCell("cell-2", 'mo.md("Hello")', null, null, hashes.moMd),
-        createNotebookCell(
-          "cell-3",
-          "x = mo.ui.slider(0, 10)",
-          null,
-          null,
-          hashes.slider,
-        ),
+        createNotebookCell("cell-1", {
+          code: "1 / 0",
+          name: null,
+          config: null,
+          code_hash: null,
+        }), // no hash
+        createNotebookCell("cell-2", {
+          code: 'mo.md("Hello")',
+          name: null,
+          config: null,
+          code_hash: hashes.moMd,
+        }),
+        createNotebookCell("cell-3", {
+          code: "x = mo.ui.slider(0, 10)",
+          name: null,
+          config: null,
+          code_hash: hashes.slider,
+        }),
       ];
 
       const session = createSession(sessionCells);
       const notebook = createNotebook(notebookCells);
-      const result = notebookStateFromSession(session, notebook);
+      const result = readExport({
+        session: session,
+        notebook: notebook,
+        filename: null,
+        appConfig: {},
+      });
 
-      expect(result).not.toBeNull();
-      invariant(result, "result is null");
+      invariant(result.snapshot, "no snapshot");
 
       // Should have notebook cell IDs in order (notebook is canonical)
-      expect(result.cellIds.inOrderIds).toEqual(
-        MultiColumn.from([["cell-1", "cell-2", "cell-3"]]).inOrderIds,
-      );
+      expect(result.snapshot.document.order).toEqual([
+        "cell-1",
+        "cell-2",
+        "cell-3",
+      ]);
 
       // Should have correct code from notebook
-      expect(result.cellData[cellId("cell-1")].code).toBe("1 / 0");
-      expect(result.cellData[cellId("cell-2")].code).toBe('mo.md("Hello")');
-      expect(result.cellData[cellId("cell-3")].code).toBe(
-        "x = mo.ui.slider(0, 10)",
-      );
+      expect(cellIn(result, "cell-1").code).toBe("1 / 0");
+      expect(cellIn(result, "cell-2").code).toBe('mo.md("Hello")');
+      expect(cellIn(result, "cell-3").code).toBe("x = mo.ui.slider(0, 10)");
 
       // cell-1: No matching session cell (hash is null), gets stub session cell
       expect(result.cellRuntime[cellId("cell-1")].output).toBeNull();