@marimo-team/frontend 0.19.3-dev8 → 0.19.4-dev0

package/src/core/export/__tests__/hooks.test.ts

@@ -0,0 +1,504 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+import { renderHook } from "@testing-library/react";
+import { createStore, Provider } from "jotai";
+import type { ReactNode } from "react";
+import * as React from "react";
+import { beforeEach, describe, expect, it, vi } from "vitest";
+import type { CellId } from "@/core/cells/ids";
+import { CellOutputId } from "@/core/cells/ids";
+import type { CellRuntimeState } from "@/core/cells/types";
+import { useEnrichCellOutputs } from "../hooks";
+
+// Mock html-to-image
+vi.mock("html-to-image", () => ({
+  toPng: vi.fn(),
+}));
+
+// Mock Logger
+vi.mock("@/utils/Logger", () => ({
+  Logger: {
+    error: vi.fn(),
+  },
+}));
+
+// Mock cellsRuntimeAtom - must be defined inline in the factory function
+vi.mock("@/core/cells/cells", async () => {
+  const { atom } = await import("jotai");
+  return {
+    cellsRuntimeAtom: atom({}),
+  };
+});
+
+import { toPng } from "html-to-image";
+import { cellsRuntimeAtom } from "@/core/cells/cells";
+import { Logger } from "@/utils/Logger";
+
+describe("useEnrichCellOutputs", () => {
+  let store: ReturnType<typeof createStore>;
+
+  beforeEach(() => {
+    vi.clearAllMocks();
+    store = createStore();
+  });
+
+  const wrapper = ({ children }: { children: ReactNode }) =>
+    React.createElement(Provider, { store }, children);
+
+  // Helper to set the mocked atom (cast to any to work around type mismatch)
+  const setCellsRuntime = (value: Record<CellId, CellRuntimeState>) => {
+    // eslint-disable-next-line @typescript-eslint/no-explicit-any
+    store.set(cellsRuntimeAtom as any, value);
+  };
+
+  const createMockCellRuntimes = (
+    cells: Record<string, Partial<CellRuntimeState>>,
+  ): Record<CellId, CellRuntimeState> => {
+    return Object.fromEntries(
+      Object.entries(cells).map(([cellId, cell]) => [
+        cellId as CellId,
+        {
+          output: cell.output || null,
+          status: cell.status || "idle",
+          interrupted: false,
+          errored: false,
+          runStartTimestamp: null,
+          runElapsedTimeMs: null,
+          stallTime: null as unknown as number,
+          ...cell,
+        } as CellRuntimeState,
+      ]),
+    ) as Record<CellId, CellRuntimeState>;
+  };
+
+  it("should return empty object when no cells need screenshots", async () => {
+    vi.spyOn(document, "getElementById");
+
+    // Set up cell runtimes with no text/html outputs
+    setCellsRuntime(
+      createMockCellRuntimes({
+        "cell-1": {
+          output: {
+            channel: "output",
+            mimetype: "text/plain",
+            data: "Hello World",
+            timestamp: 0,
+          },
+        },
+      }),
+    );
+
+    const { result } = renderHook(() => useEnrichCellOutputs(), { wrapper });
+
+    const enrichCellOutputs = result.current;
+    const output = await enrichCellOutputs();
+
+    expect(output).toEqual({});
+    expect(document.getElementById).not.toHaveBeenCalled();
+    expect(toPng).not.toHaveBeenCalled();
+  });
+
+  it("should capture screenshots for cells with text/html output", async () => {
+    const cellId = "cell-1" as CellId;
+    const mockElement = document.createElement("div");
+    const mockDataUrl = "data:image/png;base64,mockImageData";
+
+    // Mock document.getElementById
+    vi.spyOn(document, "getElementById").mockReturnValue(mockElement);
+    vi.mocked(toPng).mockResolvedValue(mockDataUrl);
+
+    setCellsRuntime(
+      createMockCellRuntimes({
+        [cellId]: {
+          output: {
+            channel: "output",
+            mimetype: "text/html",
+            data: "<div>Chart</div>",
+            timestamp: 0,
+          },
+        },
+      }),
+    );
+
+    const { result } = renderHook(() => useEnrichCellOutputs(), { wrapper });
+
+    const enrichCellOutputs = result.current;
+    const output = await enrichCellOutputs();
+
+    expect(document.getElementById).toHaveBeenCalledWith(
+      CellOutputId.create(cellId),
+    );
+    expect(toPng).toHaveBeenCalledWith(mockElement);
+    expect(output).toEqual({
+      [cellId]: ["image/png", mockDataUrl],
+    });
+  });
+
+  it("should skip cells where output has not changed", async () => {
+    const cellId = "cell-1" as CellId;
+    const mockElement = document.createElement("div");
+    const mockDataUrl = "data:image/png;base64,mockImageData";
+    const htmlData = "<div>Chart</div>";
+
+    vi.spyOn(document, "getElementById").mockReturnValue(mockElement);
+    vi.mocked(toPng).mockResolvedValue(mockDataUrl);
+
+    setCellsRuntime(
+      createMockCellRuntimes({
+        [cellId]: {
+          output: {
+            channel: "output",
+            mimetype: "text/html",
+            data: htmlData,
+            timestamp: 0,
+          },
+        },
+      }),
+    );
+
+    const { result, rerender } = renderHook(() => useEnrichCellOutputs(), {
+      wrapper,
+    });
+
+    // First call - should capture
+    let enrichCellOutputs = result.current;
+    let output = await enrichCellOutputs();
+    expect(output).toEqual({ [cellId]: ["image/png", mockDataUrl] });
+    expect(toPng).toHaveBeenCalledTimes(1);
+
+    // Rerender to get updated atom state
+    rerender();
+
+    // Second call with same output - should not capture again
+    enrichCellOutputs = result.current;
+    output = await enrichCellOutputs();
+    expect(output).toEqual({}); // Empty because output hasn't changed
+    expect(toPng).toHaveBeenCalledTimes(1); // Still only 1 call
+  });
+
+  it("should handle screenshot errors gracefully", async () => {
+    const cellId = "cell-1" as CellId;
+    const mockElement = document.createElement("div");
+    const error = new Error("Screenshot failed");
+
+    vi.spyOn(document, "getElementById").mockReturnValue(mockElement);
+    vi.mocked(toPng).mockRejectedValue(error);
+
+    setCellsRuntime(
+      createMockCellRuntimes({
+        [cellId]: {
+          output: {
+            channel: "output",
+            mimetype: "text/html",
+            data: "<div>Chart</div>",
+            timestamp: 0,
+          },
+        },
+      }),
+    );
+
+    const { result } = renderHook(() => useEnrichCellOutputs(), { wrapper });
+
+    const enrichCellOutputs = result.current;
+    const output = await enrichCellOutputs();
+
+    expect(output).toEqual({}); // Failed screenshot should be filtered out
+    expect(Logger.error).toHaveBeenCalledWith(
+      `Error screenshotting cell ${cellId}:`,
+      error,
+    );
+  });
+
+  it("should handle missing DOM elements", async () => {
+    const cellId = "cell-1" as CellId;
+
+    vi.spyOn(document, "getElementById").mockReturnValue(null);
+
+    setCellsRuntime(
+      createMockCellRuntimes({
+        [cellId]: {
+          output: {
+            channel: "output",
+            mimetype: "text/html",
+            data: "<div>Chart</div>",
+            timestamp: 0,
+          },
+        },
+      }),
+    );
+
+    const { result } = renderHook(() => useEnrichCellOutputs(), { wrapper });
+
+    const enrichCellOutputs = result.current;
+    const output = await enrichCellOutputs();
+
+    expect(output).toEqual({});
+    expect(Logger.error).toHaveBeenCalledWith(
+      `Output element not found for cell ${cellId}`,
+    );
+    expect(toPng).not.toHaveBeenCalled();
+  });
+
+  it("should process multiple cells in parallel", async () => {
+    const cell1 = "cell-1" as CellId;
+    const cell2 = "cell-2" as CellId;
+    const mockElement1 = document.createElement("div");
+    const mockElement2 = document.createElement("div");
+    const mockDataUrl1 = "data:image/png;base64,image1";
+    const mockDataUrl2 = "data:image/png;base64,image2";
+
+    vi.spyOn(document, "getElementById")
+      .mockReturnValueOnce(mockElement1)
+      .mockReturnValueOnce(mockElement2);
+
+    vi.mocked(toPng)
+      .mockResolvedValueOnce(mockDataUrl1)
+      .mockResolvedValueOnce(mockDataUrl2);
+
+    setCellsRuntime(
+      createMockCellRuntimes({
+        [cell1]: {
+          output: {
+            channel: "output",
+            mimetype: "text/html",
+            data: "<div>Chart 1</div>",
+            timestamp: 0,
+          },
+        },
+        [cell2]: {
+          output: {
+            channel: "output",
+            mimetype: "text/html",
+            data: "<div>Chart 2</div>",
+            timestamp: 0,
+          },
+        },
+      }),
+    );
+
+    const { result } = renderHook(() => useEnrichCellOutputs(), { wrapper });
+
+    const enrichCellOutputs = result.current;
+    const output = await enrichCellOutputs();
+
+    expect(output).toEqual({
+      [cell1]: ["image/png", mockDataUrl1],
+      [cell2]: ["image/png", mockDataUrl2],
+    });
+    expect(toPng).toHaveBeenCalledTimes(2);
+  });
+
+  it("should filter out null results from failed screenshots", async () => {
+    // Setup: one successful, one failed screenshot
+    const cell1 = "cell-1" as CellId;
+    const cell2 = "cell-2" as CellId;
+    const mockElement1 = document.createElement("div");
+    const mockDataUrl = "data:image/png;base64,image1";
+
+    vi.spyOn(document, "getElementById")
+      .mockReturnValueOnce(mockElement1)
+      .mockReturnValueOnce(null); // Second cell fails to find element
+
+    vi.mocked(toPng).mockResolvedValue(mockDataUrl);
+
+    setCellsRuntime(
+      createMockCellRuntimes({
+        [cell1]: {
+          output: {
+            channel: "output",
+            mimetype: "text/html",
+            data: "<div>Chart 1</div>",
+            timestamp: 0,
+          },
+        },
+        [cell2]: {
+          output: {
+            channel: "output",
+            mimetype: "text/html",
+            data: "<div>Chart 2</div>",
+            timestamp: 0,
+          },
+        },
+      }),
+    );
+
+    const { result } = renderHook(() => useEnrichCellOutputs(), { wrapper });
+
+    const enrichCellOutputs = result.current;
+    const output = await enrichCellOutputs();
+
+    // Only the successful screenshot should be in the result
+    expect(output).toEqual({
+      [cell1]: ["image/png", mockDataUrl],
+    });
+    expect(Logger.error).toHaveBeenCalledWith(
+      `Output element not found for cell ${cell2}`,
+    );
+  });
+
+  it("should only capture screenshots for cells with changed output", async () => {
+    const cellId = "cell-1" as CellId;
+    const mockElement = document.createElement("div");
+    const mockDataUrl1 = "data:image/png;base64,image1";
+    const mockDataUrl2 = "data:image/png;base64,image2";
+
+    vi.spyOn(document, "getElementById").mockReturnValue(mockElement);
+    vi.mocked(toPng)
+      .mockResolvedValueOnce(mockDataUrl1)
+      .mockResolvedValueOnce(mockDataUrl2);
+
+    // First call - cell should be captured
+    setCellsRuntime(
+      createMockCellRuntimes({
+        [cellId]: {
+          output: {
+            channel: "output",
+            mimetype: "text/html",
+            data: "<div>Chart v1</div>",
+            timestamp: 0,
+          },
+        },
+      }),
+    );
+
+    const { result, rerender } = renderHook(() => useEnrichCellOutputs(), {
+      wrapper,
+    });
+
+    // First screenshot
+    let enrichCellOutputs = result.current;
+    let output = await enrichCellOutputs();
+    expect(output).toEqual({ [cellId]: ["image/png", mockDataUrl1] });
+
+    // Second call - same output, should not be captured
+    rerender();
+    enrichCellOutputs = result.current;
+    output = await enrichCellOutputs();
+    expect(output).toEqual({});
+
+    // Third call - output changed, should be captured
+    setCellsRuntime(
+      createMockCellRuntimes({
+        [cellId]: {
+          output: {
+            channel: "output",
+            mimetype: "text/html",
+            data: "<div>Chart v2</div>", // Changed!
+            timestamp: 0,
+          },
+        },
+      }),
+    );
+
+    rerender();
+    enrichCellOutputs = result.current;
+    output = await enrichCellOutputs();
+    expect(output).toEqual({ [cellId]: ["image/png", mockDataUrl2] });
+    expect(toPng).toHaveBeenCalledTimes(2);
+  });
+
+  it("should ignore cells with non-text/html mimetype", async () => {
+    vi.spyOn(document, "getElementById");
+
+    setCellsRuntime(
+      createMockCellRuntimes({
+        "cell-1": {
+          output: {
+            channel: "output",
+            mimetype: "application/json",
+            data: '{"key": "value"}',
+            timestamp: 0,
+          },
+        },
+        "cell-2": {
+          output: {
+            channel: "output",
+            mimetype: "text/plain",
+            data: "Plain text",
+            timestamp: 0,
+          },
+        },
+        "cell-3": {
+          output: {
+            channel: "output",
+            mimetype: "image/png",
+            data: "data:image/png;base64,existing",
+            timestamp: 0,
+          },
+        },
+      }),
+    );
+
+    const { result } = renderHook(() => useEnrichCellOutputs(), { wrapper });
+
+    const enrichCellOutputs = result.current;
+    const output = await enrichCellOutputs();
+
+    // None of these should trigger screenshots
+    expect(output).toEqual({});
+    expect(document.getElementById).not.toHaveBeenCalled();
+    expect(toPng).not.toHaveBeenCalled();
+  });
+
+  it("should ignore cells with null or undefined output", async () => {
+    vi.spyOn(document, "getElementById");
+
+    setCellsRuntime(
+      createMockCellRuntimes({
+        "cell-1": {
+          output: null,
+        },
+        "cell-2": {
+          output: undefined,
+        },
+      }),
+    );
+
+    const { result } = renderHook(() => useEnrichCellOutputs(), { wrapper });
+
+    const enrichCellOutputs = result.current;
+    const output = await enrichCellOutputs();
+
+    expect(output).toEqual({});
+    expect(document.getElementById).not.toHaveBeenCalled();
+    expect(toPng).not.toHaveBeenCalled();
+  });
+
+  it("should return correctly formatted result with CellId and tuple", async () => {
+    // Expected format: Record<CellId, ["image/png", string]>
+    const cellId = "test-cell" as CellId;
+    const mockElement = document.createElement("div");
+    const mockDataUrl = "data:image/png;base64,testData";
+
+    vi.spyOn(document, "getElementById").mockReturnValue(mockElement);
+    vi.mocked(toPng).mockResolvedValue(mockDataUrl);
+
+    setCellsRuntime(
+      createMockCellRuntimes({
+        [cellId]: {
+          output: {
+            channel: "output",
+            mimetype: "text/html",
+            data: "<div>Content</div>",
+            timestamp: 0,
+          },
+        },
+      }),
+    );
+
+    const { result } = renderHook(() => useEnrichCellOutputs(), { wrapper });
+
+    const enrichCellOutputs = result.current;
+    const output = await enrichCellOutputs();
+
+    // Verify the exact return type structure
+    expect(output).toHaveProperty(cellId);
+    const cellOutput = output[cellId];
+    expect(cellOutput).toBeDefined();
+    expect(Array.isArray(cellOutput)).toBe(true);
+    if (cellOutput) {
+      expect(cellOutput[0]).toBe("image/png");
+      expect(cellOutput[1]).toBe(mockDataUrl);
+    }
+  });
+});

package/src/core/export/hooks.ts

@@ -1,7 +1,12 @@
 /* Copyright 2026 Marimo. All rights reserved. */
-import { useAtomValue } from "jotai";
+import { toPng } from "html-to-image";
+import { atom, useAtom, useAtomValue } from "jotai";
 import { appConfigAtom } from "@/core/config/config";
 import { useInterval } from "@/hooks/useInterval";
+import { Logger } from "@/utils/Logger";
+import { Objects } from "@/utils/objects";
+import { cellsRuntimeAtom } from "../cells/cells";
+import { type CellId, CellOutputId } from "../cells/ids";
 import { connectionAtom } from "../network/connection";
 import { useRequestClient } from "../network/requests";
 import { VirtualFileTracker } from "../static/virtual-file-tracker";
@@ -21,8 +26,13 @@ export function useAutoExport() {
   const markdownDisabled = !markdownEnabled || !isConnected;
   const htmlDisabled = !htmlEnabled || !isConnected;
   const ipynbDisabled = !ipynbEnabled || !isConnected;
-  const { autoExportAsHTML, autoExportAsIPYNB, autoExportAsMarkdown } =
-    useRequestClient();
+  const {
+    autoExportAsHTML,
+    autoExportAsIPYNB,
+    autoExportAsMarkdown,
+    updateCellOutputs,
+  } = useRequestClient();
+  const takeScreenshots = useEnrichCellOutputs();
 
   useInterval(
     async () => {
@@ -50,12 +60,91 @@ export function useAutoExport() {
 
   useInterval(
     async () => {
+      const cellsToOutput = await takeScreenshots();
+      if (Object.keys(cellsToOutput).length > 0) {
+        await updateCellOutputs({
+          cellIdsToOutput: cellsToOutput,
+        });
+      }
       await autoExportAsIPYNB({
         download: false,
       });
     },
     // Run every 5 seconds, or when the document becomes visible
     // Ignore if the document is not visible
-    { delayMs: DELAY, whenVisible: true, disabled: ipynbDisabled },
+    // Skip if running to ensure no race conditions between screenshot and export
+    {
+      delayMs: DELAY,
+      whenVisible: true,
+      disabled: ipynbDisabled,
+      skipIfRunning: true,
+    },
   );
 }
+
+// We track cells that need screenshots, these will be exported to IPYNB
+const richCellsToOutputAtom = atom<Record<CellId, unknown>>({});
+
+/**
+ * Take screenshots of cells with HTML outputs. These images will be sent to the backend to be exported to IPYNB.
+ * @returns A map of cell IDs to their screenshots data.
+ */
+export function useEnrichCellOutputs() {
+  const [richCellsOutput, setRichCellsOutput] = useAtom(richCellsToOutputAtom);
+  const cellRuntimes = useAtomValue(cellsRuntimeAtom);
+
+  return async (): Promise<Record<CellId, ["image/png", string]>> => {
+    const trackedCellsOutput: Record<CellId, unknown> = {};
+
+    const cellsToCaptureScreenshot: [CellId, unknown][] = [];
+    for (const [cellId, runtime] of Objects.entries(cellRuntimes)) {
+      const outputData = runtime.output?.data;
+      const outputHasChanged = richCellsOutput[cellId] !== outputData;
+      // Track latest output for this cell
+      trackedCellsOutput[cellId] = outputData;
+      if (
+        runtime.output?.mimetype === "text/html" &&
+        outputData &&
+        outputHasChanged
+      ) {
+        cellsToCaptureScreenshot.push([cellId, runtime]);
+      }
+    }
+    // Always update tracked outputs, this ensures data is fresh for the next run
+    setRichCellsOutput(trackedCellsOutput);
+
+    if (cellsToCaptureScreenshot.length === 0) {
+      return {};
+    }
+
+    // Capture screenshots
+    const results = await Promise.all(
+      cellsToCaptureScreenshot.map(async ([cellId]) => {
+        const outputElement = document.getElementById(
+          CellOutputId.create(cellId),
+        );
+        if (!outputElement) {
+          Logger.error(`Output element not found for cell ${cellId}`);
+          return null;
+        }
+
+        try {
+          const dataUrl = await toPng(outputElement);
+          return [cellId, ["image/png", dataUrl]] as [
+            CellId,
+            ["image/png", string],
+          ];
+        } catch (error) {
+          Logger.error(`Error screenshotting cell ${cellId}:`, error);
+          return null;
+        }
+      }),
+    );
+
+    return Objects.fromEntries(
+      results.filter(
+        (result): result is [CellId, ["image/png", string]] => result !== null,
+      ),
+    );
+  };
+}

package/src/core/islands/bridge.ts

@@ -171,6 +171,7 @@ export class IslandsPyodideBridge implements RunRequests, EditRequests {
   autoExportAsHTML = throwNotImplemented;
   autoExportAsMarkdown = throwNotImplemented;
   autoExportAsIPYNB = throwNotImplemented;
+  updateCellOutputs = throwNotImplemented;
   addPackage = throwNotImplemented;
   removePackage = throwNotImplemented;
   getPackageList = throwNotImplemented;

package/src/core/kernel/__tests__/handlers.test.ts

@@ -142,8 +142,8 @@ describe("buildCellData", () => {
         cell2: "y = 2",
       },
       last_execution_time: {
-        cell1: 1234567890,
-        cell2: 1234567891,
+        cell1: 1_234_567_890,
+        cell2: 1_234_567_891,
       },
       app_config: {
         width: "normal",

package/src/core/kernel/state.ts

@@ -1,3 +1,4 @@
+/* Copyright 2026 Marimo. All rights reserved. */
 import { atom } from "jotai";
 import { waitFor } from "../state/jotai";
 

package/src/core/network/__tests__/requests-lazy.test.ts

@@ -1,4 +1,4 @@
-/* Copyright 2024 Marimo. All rights reserved. */
+/* Copyright 2026 Marimo. All rights reserved. */
 
 import { beforeEach, describe, expect, it, vi } from "vitest";
 import type { RuntimeManager } from "../../runtime/runtime";

package/src/core/network/__tests__/requests-network.test.ts

@@ -34,12 +34,6 @@ vi.mock("../connection", () => ({
   waitForConnectionOpen: vi.fn().mockResolvedValue(undefined),
 }));
 
-vi.mock("../../state/jotai", () => ({
-  store: {
-    get: vi.fn(() => true),
-  },
-}));
-
 describe("createNetworkRequests", () => {
   let mockClient: any;
   let capturedCalls: Map<string, { hasParams: boolean; endpoint: string }>;
@@ -80,18 +74,6 @@ describe("createNetworkRequests", () => {
   });
 
   describe("special behavior", () => {
-    it("sendRun should drop requests if not connected", async () => {
-      const { store } = await import("../../state/jotai");
-
-      vi.mocked(store.get).mockReturnValue(false);
-
-      const requests = createNetworkRequests();
-      const result = await requests.sendRun({} as any);
-
-      expect(result).toBe(null);
-      expect(mockClient.POST).not.toHaveBeenCalled();
-    });
-
     it("exportAsHTML should set assetUrl in dev/test mode", async () => {
       const originalEnv = process.env.NODE_ENV;
       process.env.NODE_ENV = "development";