@marimo-team/frontend 0.19.7-dev34 → 0.19.7-dev36

package/src/components/editor/renderers/vertical-layout/vertical-layout.tsx

@@ -41,7 +41,11 @@ import { downloadAsHTML } from "@/core/static/download-html";
 import { isStaticNotebook } from "@/core/static/static-state";
 import { isWasm } from "@/core/wasm/utils";
 import { cn } from "@/utils/cn";
-import { downloadBlob, downloadHTMLAsImage } from "@/utils/download";
+import {
+  ADD_PRINTING_CLASS,
+  downloadBlob,
+  downloadHTMLAsImage,
+} from "@/utils/download";
 import { Filenames } from "@/utils/filenames";
 import { FloatingOutline } from "../../chrome/panels/outline/floating-outline";
 import { cellDomProps } from "../../common";
@@ -185,7 +189,12 @@ const ActionButtons: React.FC<{
     if (!app) {
       return;
     }
-    await downloadHTMLAsImage({ element: app, filename: document.title });
+    await downloadHTMLAsImage({
+      element: app,
+      filename: document.title,
+      // Add body.printing ONLY when converting the whole notebook to a screenshot
+      prepare: ADD_PRINTING_CLASS,
+    });
   };
 
   const handleDownloadAsHTML = async () => {
@@ -271,7 +280,7 @@ const ActionButtons: React.FC<{
     <div
       data-testid="notebook-actions-dropdown"
       className={cn(
-        "right-0 top-0 z-50 m-4 no-print flex gap-2 print:hidden",
+        "right-0 top-0 z-50 m-4 print:hidden flex gap-2",
         // If the notebook is static, we have a banner at the top, so
         // we can't use fixed positioning. Ideally this is sticky, but the
         // current dom structure makes that difficult.
@@ -284,7 +293,7 @@ const ActionButtons: React.FC<{
             <MoreHorizontalIcon className="w-4 h-4" />
           </Button>
         </DropdownMenuTrigger>
-        <DropdownMenuContent align="end" className="no-print w-[220px]">
+        <DropdownMenuContent align="end" className="print:hidden w-[220px]">
           {actions}
         </DropdownMenuContent>
       </DropdownMenu>

package/src/components/home/components.tsx

@@ -80,7 +80,7 @@ export const OpenTutorialDropDown: React.FC = () => {
           <CaretDownIcon className="w-3 h-3 ml-1" />
         </Button>
       </DropdownMenuTrigger>
-      <DropdownMenuContent side="bottom" align="end" className="no-print">
+      <DropdownMenuContent side="bottom" align="end" className="print:hidden">
         {Objects.entries(TUTORIALS).map(
           ([tutorialId, [label, Icon, description]]) => (
             <DropdownMenuItem

package/src/components/static-html/static-banner.tsx

@@ -36,7 +36,7 @@ export const StaticBanner: React.FC = () => {
 
   return (
     <div
-      className="px-4 py-2 bg-(--sky-2) border-b border-(--sky-7) text-(--sky-11) flex justify-between items-center gap-4 no-print text-sm"
+      className="px-4 py-2 bg-(--sky-2) border-b border-(--sky-7) text-(--sky-11) flex justify-between items-center gap-4 print:hidden text-sm"
       data-testid="static-notebook-banner"
     >
       <span>

package/src/components/ui/dropdown-menu.tsx

@@ -83,7 +83,7 @@ const DropdownMenuContent = React.forwardRef<
         sideOffset={sideOffset}
         className={cn(
           menuContentCommon(),
-          "animate-in data-[side=bottom]:slide-in-from-top-2 data-[side=left]:slide-in-from-right-2 data-[side=right]:slide-in-from-left-2 data-[side=top]:slide-in-from-bottom-2",
+          "animate-in data-[side=bottom]:slide-in-from-top-2 data-[side=left]:slide-in-from-right-2 data-[side=right]:slide-in-from-left-2 data-[side=top]:slide-in-from-bottom-2 print:hidden",
           scrollable && "overflow-auto",
           className,
         )}

package/src/components/ui/table.tsx

@@ -7,7 +7,7 @@ const Table = React.forwardRef<
   HTMLTableElement,
   React.HTMLAttributes<HTMLTableElement>
 >(({ className, ...props }, ref) => (
-  <div className="w-full overflow-auto scrollbar-thin flex-1">
+  <div className="w-full overflow-auto scrollbar-thin flex-1 print:overflow-hidden">
     <table
       ref={ref}
       className={cn("w-full caption-bottom text-sm", className)}

package/src/core/config/feature-flag.tsx

@@ -25,7 +25,7 @@ const defaultValues: ExperimentalFeatures = {
   chat_modes: false,
   cache_panel: false,
   external_agents: import.meta.env.DEV,
-  server_side_pdf_export: false,
+  server_side_pdf_export: true,
 };
 
 export function getFeatureFlag<T extends keyof ExperimentalFeatures>(

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

@@ -10,6 +10,7 @@ import { CellOutputId } from "@/core/cells/ids";
 import type { CellRuntimeState } from "@/core/cells/types";
 import { ProgressState } from "@/utils/progress";
 import {
+  captureTracker,
   updateCellOutputsWithScreenshots,
   useEnrichCellOutputs,
 } from "../hooks";
@@ -52,6 +53,7 @@ describe("useEnrichCellOutputs", () => {
   beforeEach(() => {
     vi.clearAllMocks();
     store = createStore();
+    captureTracker.reset();
   });
 
   const wrapper = ({ children }: { children: ReactNode }) =>
@@ -103,7 +105,7 @@ describe("useEnrichCellOutputs", () => {
     const { result } = renderHook(() => useEnrichCellOutputs(), { wrapper });
 
     const takeScreenshots = result.current;
-    const output = await takeScreenshots({ progress, snappy: false });
+    const output = await takeScreenshots({ progress });
 
     expect(output).toEqual({});
     expect(document.getElementById).not.toHaveBeenCalled();
@@ -135,7 +137,7 @@ describe("useEnrichCellOutputs", () => {
     const { result } = renderHook(() => useEnrichCellOutputs(), { wrapper });
 
     const takeScreenshots = result.current;
-    const output = await takeScreenshots({ progress, snappy: false });
+    const output = await takeScreenshots({ progress });
 
     expect(document.getElementById).toHaveBeenCalledWith(
       CellOutputId.create(cellId),
@@ -152,50 +154,6 @@ describe("useEnrichCellOutputs", () => {
     });
   });
 
-  it("should pass snappy=true to toPng with includeStyleProperties", 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 takeScreenshots = result.current;
-    const output = await takeScreenshots({ progress, snappy: true });
-
-    expect(document.getElementById).toHaveBeenCalledWith(
-      CellOutputId.create(cellId),
-    );
-    // When snappy=true, includeStyleProperties should be set
-    expect(toPng).toHaveBeenCalledWith(
-      mockElement,
-      expect.objectContaining({
-        filter: expect.any(Function),
-        onImageErrorHandler: expect.any(Function),
-        includeStyleProperties: expect.any(Array),
-      }),
-    );
-    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");
@@ -224,7 +182,7 @@ describe("useEnrichCellOutputs", () => {
 
     // First call - should capture
     let takeScreenshots = result.current;
-    let output = await takeScreenshots({ progress, snappy: false });
+    let output = await takeScreenshots({ progress });
     expect(output).toEqual({ [cellId]: ["image/png", mockDataUrl] });
     expect(toPng).toHaveBeenCalledTimes(1);
 
@@ -233,7 +191,7 @@ describe("useEnrichCellOutputs", () => {
 
     // Second call with same output - should not capture again
     takeScreenshots = result.current;
-    output = await takeScreenshots({ progress, snappy: false });
+    output = await takeScreenshots({ progress });
     expect(output).toEqual({}); // Empty because output hasn't changed
     expect(toPng).toHaveBeenCalledTimes(1); // Still only 1 call
   });
@@ -262,7 +220,7 @@ describe("useEnrichCellOutputs", () => {
     const { result } = renderHook(() => useEnrichCellOutputs(), { wrapper });
 
     const takeScreenshots = result.current;
-    const output = await takeScreenshots({ progress, snappy: false });
+    const output = await takeScreenshots({ progress });
 
     expect(output).toEqual({}); // Failed screenshot should be filtered out
     expect(Logger.error).toHaveBeenCalledWith(
@@ -271,6 +229,50 @@ describe("useEnrichCellOutputs", () => {
     );
   });
 
+  it("should retry failed screenshots on next call", async () => {
+    const cellId = "cell-1" as CellId;
+    const mockElement = document.createElement("div");
+    const error = new Error("Screenshot failed");
+    const mockDataUrl = "data:image/png;base64,retrySuccess";
+
+    vi.spyOn(document, "getElementById").mockReturnValue(mockElement);
+    // First call fails, second call succeeds
+    vi.mocked(toPng)
+      .mockRejectedValueOnce(error)
+      .mockResolvedValueOnce(mockDataUrl);
+
+    setCellsRuntime(
+      createMockCellRuntimes({
+        [cellId]: {
+          output: {
+            channel: "output",
+            mimetype: "text/html",
+            data: "<div>Chart</div>",
+            timestamp: 0,
+          },
+        },
+      }),
+    );
+
+    const { result, rerender } = renderHook(() => useEnrichCellOutputs(), {
+      wrapper,
+    });
+
+    // First call - screenshot fails
+    let takeScreenshots = result.current;
+    let output = await takeScreenshots({ progress });
+    expect(output).toEqual({});
+    expect(Logger.error).toHaveBeenCalled();
+
+    rerender();
+
+    // Second call - should retry since the first one failed
+    takeScreenshots = result.current;
+    output = await takeScreenshots({ progress });
+    expect(output).toEqual({ [cellId]: ["image/png", mockDataUrl] });
+    expect(toPng).toHaveBeenCalledTimes(2);
+  });
+
   it("should handle missing DOM elements", async () => {
     const cellId = "cell-1" as CellId;
 
@@ -292,7 +294,7 @@ describe("useEnrichCellOutputs", () => {
     const { result } = renderHook(() => useEnrichCellOutputs(), { wrapper });
 
     const takeScreenshots = result.current;
-    const output = await takeScreenshots({ progress, snappy: false });
+    const output = await takeScreenshots({ progress });
 
     expect(output).toEqual({});
     expect(Logger.error).toHaveBeenCalledWith(
@@ -341,7 +343,7 @@ describe("useEnrichCellOutputs", () => {
     const { result } = renderHook(() => useEnrichCellOutputs(), { wrapper });
 
     const takeScreenshots = result.current;
-    const output = await takeScreenshots({ progress, snappy: false });
+    const output = await takeScreenshots({ progress });
 
     expect(output).toEqual({
       [cell1]: ["image/png", mockDataUrl1],
@@ -387,7 +389,7 @@ describe("useEnrichCellOutputs", () => {
     const { result } = renderHook(() => useEnrichCellOutputs(), { wrapper });
 
     const takeScreenshots = result.current;
-    const output = await takeScreenshots({ progress, snappy: false });
+    const output = await takeScreenshots({ progress });
 
     // Only the successful screenshot should be in the result
     expect(output).toEqual({
@@ -429,13 +431,13 @@ describe("useEnrichCellOutputs", () => {
 
     // First screenshot
     let takeScreenshots = result.current;
-    let output = await takeScreenshots({ progress, snappy: false });
+    let output = await takeScreenshots({ progress });
     expect(output).toEqual({ [cellId]: ["image/png", mockDataUrl1] });
 
     // Second call - same output, should not be captured
     rerender();
     takeScreenshots = result.current;
-    output = await takeScreenshots({ progress, snappy: false });
+    output = await takeScreenshots({ progress });
     expect(output).toEqual({});
 
     // Third call - output changed, should be captured
@@ -454,7 +456,7 @@ describe("useEnrichCellOutputs", () => {
 
     rerender();
     takeScreenshots = result.current;
-    output = await takeScreenshots({ progress, snappy: false });
+    output = await takeScreenshots({ progress });
     expect(output).toEqual({ [cellId]: ["image/png", mockDataUrl2] });
     expect(toPng).toHaveBeenCalledTimes(2);
   });
@@ -494,7 +496,7 @@ describe("useEnrichCellOutputs", () => {
     const { result } = renderHook(() => useEnrichCellOutputs(), { wrapper });
 
     const takeScreenshots = result.current;
-    const output = await takeScreenshots({ progress, snappy: false });
+    const output = await takeScreenshots({ progress });
 
     // None of these should trigger screenshots
     expect(output).toEqual({});
@@ -519,7 +521,7 @@ describe("useEnrichCellOutputs", () => {
     const { result } = renderHook(() => useEnrichCellOutputs(), { wrapper });
 
     const takeScreenshots = result.current;
-    const output = await takeScreenshots({ progress, snappy: false });
+    const output = await takeScreenshots({ progress });
 
     expect(output).toEqual({});
     expect(document.getElementById).not.toHaveBeenCalled();
@@ -551,7 +553,7 @@ describe("useEnrichCellOutputs", () => {
     const { result } = renderHook(() => useEnrichCellOutputs(), { wrapper });
 
     const takeScreenshots = result.current;
-    const output = await takeScreenshots({ progress, snappy: false });
+    const output = await takeScreenshots({ progress });
 
     // Verify the exact return type structure
     expect(output).toHaveProperty(cellId);

package/src/core/export/hooks.ts

@@ -1,9 +1,10 @@
 /* Copyright 2026 Marimo. All rights reserved. */
-import { atom, useAtom, useAtomValue } from "jotai";
+import { useAtomValue } from "jotai";
 import type { MimeType } from "@/components/editor/Output";
 import { toast } from "@/components/ui/use-toast";
 import { appConfigAtom } from "@/core/config/config";
 import { useInterval } from "@/hooks/useInterval";
+import { AsyncCaptureTracker } from "@/utils/async-capture-tracker";
 import { getImageDataUrlForCell } from "@/utils/download";
 import { Logger } from "@/utils/Logger";
 import { Objects } from "@/utils/objects";
@@ -67,7 +68,6 @@ export function useAutoExport() {
       const screenshotFn = () =>
         takeScreenshots({
           progress: ProgressState.indeterminate(),
-          snappy: true,
         });
       await updateCellOutputsWithScreenshots({
         takeScreenshots: screenshotFn,
@@ -89,9 +89,6 @@ export function useAutoExport() {
   );
 }
 
-// We track cells that need screenshots, these will be exported to IPYNB
-const richCellsToOutputAtom = atom<Record<CellId, unknown>>({});
-
 // MIME types to capture screenshots for
 const MIME_TYPES_TO_CAPTURE_SCREENSHOTS = new Set<MimeType>([
   "text/html",
@@ -101,66 +98,109 @@ const MIME_TYPES_TO_CAPTURE_SCREENSHOTS = new Set<MimeType>([
   "application/vnd.vega.v6+json",
 ]);
 
-type ScreenshotResults = Record<CellId, ["image/png", string]>;
+type ScreenshotResult = ["image/png", string];
+type ScreenshotResults = Record<CellId, ScreenshotResult>;
+
+// Only marks cells as captured after successful screenshot.
+export const captureTracker = new AsyncCaptureTracker<
+  CellId,
+  ScreenshotResult
+>();
+
+interface UseEnrichCellOutputsOptions {
+  progress: ProgressState;
+}
 
 /**
  * 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);
+export function useEnrichCellOutputs(): (
+  opts: UseEnrichCellOutputsOptions,
+) => Promise<ScreenshotResults> {
   const cellRuntimes = useAtomValue(cellsRuntimeAtom);
 
-  return async ({
-    progress,
-    snappy,
-  }: {
-    progress: ProgressState;
-    snappy: boolean;
-  }): Promise<ScreenshotResults> => {
-    const trackedCellsOutput: Record<CellId, unknown> = {};
+  return async (
+    opts: UseEnrichCellOutputsOptions,
+  ): Promise<ScreenshotResults> => {
+    const { progress } = opts;
+
+    // Prune tracked state for cells that no longer exist
+    const currentCellIds = new Set(Objects.keys(cellRuntimes));
+    captureTracker.prune(currentCellIds);
 
     const cellsToCaptureScreenshot: [CellId, unknown][] = [];
+    const inFlightWaiters: {
+      cellId: CellId;
+      promise: Promise<ScreenshotResult | undefined>;
+    }[] = [];
+
     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 &&
         MIME_TYPES_TO_CAPTURE_SCREENSHOTS.has(runtime.output.mimetype) &&
-        outputData &&
-        outputHasChanged
+        outputData
       ) {
-        cellsToCaptureScreenshot.push([cellId, runtime]);
+        if (captureTracker.needsCapture(cellId, outputData)) {
+          cellsToCaptureScreenshot.push([cellId, outputData]);
+        } else {
+          // If already in-flight with the same value, await its result
+          const promise = captureTracker.waitForInFlight(cellId, outputData);
+          if (promise) {
+            inFlightWaiters.push({ cellId, promise });
+          }
+        }
       }
     }
-    // Always update tracked outputs, this ensures data is fresh for the next run
-    setRichCellsOutput(trackedCellsOutput);
 
-    if (cellsToCaptureScreenshot.length === 0) {
+    if (cellsToCaptureScreenshot.length === 0 && inFlightWaiters.length === 0) {
       return {};
     }
 
-    // Capture screenshots
-    const total = cellsToCaptureScreenshot.length;
-    progress.addTotal(total);
+    // Start the progress bar for new captures only
+    if (cellsToCaptureScreenshot.length > 0) {
+      progress.addTotal(cellsToCaptureScreenshot.length);
+    }
+
+    // Capture screenshots — each key gets its own AbortSignal so
+    // aborting one cell does not affect the others.
     const results: ScreenshotResults = {};
-    for (const [cellId] of cellsToCaptureScreenshot) {
+    for (const [cellId, outputData] of cellsToCaptureScreenshot) {
+      const handle = captureTracker.startCapture(cellId, outputData);
       try {
-        const dataUrl = await getImageDataUrlForCell(cellId, snappy);
+        const dataUrl = await getImageDataUrlForCell(cellId);
+        if (handle.signal.aborted) {
+          continue;
+        }
         if (!dataUrl) {
           Logger.error(`Failed to capture screenshot for cell ${cellId}`);
+          handle.markFailed();
           continue;
         }
-        results[cellId] = ["image/png", dataUrl];
+        const result: ScreenshotResult = ["image/png", dataUrl];
+        results[cellId] = result;
+        handle.markCaptured(result);
       } catch (error) {
         Logger.error(`Error screenshotting cell ${cellId}:`, error);
+        handle.markFailed();
       } finally {
         progress.increment(1);
       }
     }
 
+    // Await in-flight captures started by concurrent callers
+    const settled = await Promise.allSettled(
+      inFlightWaiters.map(({ promise }) => promise),
+    );
+    for (const [i, { cellId }] of inFlightWaiters.entries()) {
+      const result =
+        settled[i].status === "fulfilled" ? settled[i].value : undefined;
+      if (result) {
+        results[cellId] = result;
+      }
+    }
+
     return results;
   };
 }

package/src/css/app/print.css

@@ -1,24 +1,10 @@
 @reference "../globals.css";
 
-/* Hide on print
-To hide an element on print, add the class `no-print` to it.
-When printing, the class `printing` is added to the body element to enable this rule. */
-body.printing .no-print {
-  display: none !important;
-}
-
 body.printing #App {
   /* Full screen print */
   height: fit-content !important;
 }
 
-/* When printing the output of a cell, this unset the max-height set by the notebook to capture the full output */
-.printing-output {
-  max-height: none !important;
-
-  @apply bg-background;
-}
-
 @media print {
   * {
     -webkit-print-color-adjust: exact;