@marimo-team/islands 0.19.7-dev24 → 0.19.7-dev26

package/dist/{types-BIvEidLq.js → types-ZLLMdAtn.js}

@@ -5,15 +5,15 @@ var _a;
 import { s as __toESM } from "./chunk-BNovOVIE.js";
 import { t as require_react } from "./react-DdA8EBol.js";
 import { t as require_compiler_runtime } from "./compiler-runtime-DHFVbq0b.js";
-import { D as dispatchDiscreteCustomEvent, E as Primitive, M as composeEventHandlers, N as createLucideIcon, O as createSlot, S as useId, T as useCallbackRef, _ as Arrow, b as createPopperScope, c as withSmartCollisionBoundary, g as Anchor, h as Portal, i as useFocusGuards, j as createContextScope, m as Presence, n as hideOthers, o as MAX_HEIGHT_OFFSET, p as useControllableState, r as FocusScope, s as withFullScreenAsRoot, t as Combination_default, u as StyleNamespace, v as Content, w as DismissableLayer, y as Root2$1 } from "./Combination-BOmAhdTT.js";
-import { t as Check } from "./check-S8ldILuD.js";
-import { B as ChevronDown, C as useDirection, L as $b5e257d569688ac6$export$619500959fc48b26, M as $a916eb452884faea$export$b7a616150fdb9f44, N as $488c6ddbf4ef74c2$export$cc77c4ff7e8673c5, P as $18f2051aff69b9bf$export$43bb16f9c6d9e3f7, R as X, _ as menuItemVariants, b as menuSubTriggerVariants, g as menuControlVariants, h as menuControlCheckVariants, j as upperFirst_default, m as menuContentCommon, v as menuLabelVariants, w as createCollection, y as menuSeparatorVariants, z as ChevronUp } from "./label-o68OaAQk.js";
+import { D as dispatchDiscreteCustomEvent, E as Primitive, M as composeEventHandlers, N as createLucideIcon, O as createSlot, S as useId, T as useCallbackRef, _ as Arrow, b as createPopperScope, c as withSmartCollisionBoundary, g as Anchor, h as Portal, i as useFocusGuards, j as createContextScope, m as Presence, n as hideOthers, o as MAX_HEIGHT_OFFSET, p as useControllableState, r as FocusScope, s as withFullScreenAsRoot, t as Combination_default, u as StyleNamespace, v as Content, w as DismissableLayer, y as Root2$1 } from "./Combination-Bg-xN8JV.js";
+import { t as Check } from "./check-CM_kewwn.js";
+import { B as ChevronDown, C as useDirection, L as $b5e257d569688ac6$export$619500959fc48b26, M as $a916eb452884faea$export$b7a616150fdb9f44, N as $488c6ddbf4ef74c2$export$cc77c4ff7e8673c5, P as $18f2051aff69b9bf$export$43bb16f9c6d9e3f7, R as X, _ as menuItemVariants, b as menuSubTriggerVariants, g as menuControlVariants, h as menuControlCheckVariants, j as upperFirst_default, m as menuContentCommon, v as menuLabelVariants, w as createCollection, y as menuSeparatorVariants, z as ChevronUp } from "./label-DC2pbeUJ.js";
 import { n as clsx_default } from "./clsx-D2KVTYnW.js";
-import { c as useComposedRefs, l as Events, s as composeRefs, t as Button, u as cn } from "./button-D6ZIdUA3.js";
-import { s as Logger } from "./hotkeys-C4e3s3sJ.js";
+import { c as useComposedRefs, l as Events, s as composeRefs, t as Button, u as cn } from "./button-BWvsJ2Wr.js";
+import { s as Logger } from "./hotkeys-B5WnGZXF.js";
 import { t as require_jsx_runtime } from "./jsx-runtime-CTBg5pdT.js";
 import { t as require_react_dom } from "./react-dom-DJW8xUDg.js";
-import { h as useEvent_default } from "./useTheme-D0zT6HZe.js";
+import { h as useEvent_default } from "./useTheme-BYG2SH8J.js";
 import { t as toString_default } from "./toString-C4TLO6FA.js";
 import { i as debounce_default, n as Constants } from "./constants-DrOu5vvd.js";
 import { t as memoizeLastValue } from "./once-BqS42WgZ.js";

package/dist/{useAsyncData-DFBEii2l.js → useAsyncData-BjNwqCfS.js}

@@ -1,7 +1,7 @@
 import { s as __toESM } from "./chunk-BNovOVIE.js";
 import { t as require_react } from "./react-DdA8EBol.js";
 import { t as require_compiler_runtime } from "./compiler-runtime-DHFVbq0b.js";
-import { h as useEvent_default } from "./useTheme-D0zT6HZe.js";
+import { h as useEvent_default } from "./useTheme-BYG2SH8J.js";
 import { t as invariant } from "./invariant-D9QLJ4SZ.js";
 var import_compiler_runtime = require_compiler_runtime(), import_react = /* @__PURE__ */ __toESM(require_react(), 1), Result = {
   error(e, s) {

package/dist/{useDeepCompareMemoize-HKwezUBx.js → useDeepCompareMemoize-CfoxVor3.js}

@@ -4,13 +4,13 @@ var __publicField = (obj, key, value) => __defNormalProp(obj, typeof key !== "sy
 import { s as __toESM, t as __commonJSMin } from "./chunk-BNovOVIE.js";
 import { t as require_react } from "./react-DdA8EBol.js";
 import { t as require_compiler_runtime } from "./compiler-runtime-DHFVbq0b.js";
-import { N as createLucideIcon } from "./Combination-BOmAhdTT.js";
-import { a as cva, u as cn } from "./button-D6ZIdUA3.js";
-import { s as Logger } from "./hotkeys-C4e3s3sJ.js";
+import { N as createLucideIcon } from "./Combination-Bg-xN8JV.js";
+import { a as cva, u as cn } from "./button-BWvsJ2Wr.js";
+import { s as Logger } from "./hotkeys-B5WnGZXF.js";
 import { t as require_jsx_runtime } from "./jsx-runtime-CTBg5pdT.js";
-import { f as waitFor, p as isIslands, u as store, x as atom } from "./useTheme-D0zT6HZe.js";
+import { f as waitFor, p as isIslands, u as store, x as atom } from "./useTheme-BYG2SH8J.js";
 import { r as KnownQueryParams } from "./constants-DrOu5vvd.js";
-import { i as tableFromIPC } from "./loader-WBQoIunJ.js";
+import { i as tableFromIPC } from "./loader-B0KEFFi-.js";
 var CircleQuestionMark = createLucideIcon("circle-question-mark", [
   ["circle", {
     cx: "12",

package/dist/{useIframeCapabilities-wWXCXV88.js → useIframeCapabilities-BBO_R0ww.js}

@@ -1,7 +1,7 @@
 import { s as __toESM } from "./chunk-BNovOVIE.js";
 import { t as require_react } from "./react-DdA8EBol.js";
 import { t as require_compiler_runtime } from "./compiler-runtime-DHFVbq0b.js";
-import { s as Logger } from "./hotkeys-C4e3s3sJ.js";
+import { s as Logger } from "./hotkeys-B5WnGZXF.js";
 import { n as once } from "./once-BqS42WgZ.js";
 function testStorage(e) {
   try {

package/dist/{useTheme-D0zT6HZe.js → useTheme-BYG2SH8J.js}

@@ -1,7 +1,7 @@
 import { s as __toESM } from "./chunk-BNovOVIE.js";
 import { t as require_react } from "./react-DdA8EBol.js";
 import { t as require_compiler_runtime } from "./compiler-runtime-DHFVbq0b.js";
-import { a as resolvePlatform, n as OverridingHotkeyProvider, s as Logger } from "./hotkeys-C4e3s3sJ.js";
+import { a as resolvePlatform, n as OverridingHotkeyProvider, s as Logger } from "./hotkeys-B5WnGZXF.js";
 import { t as _baseIsEqual_default } from "./_baseIsEqual-CBSjxu-D.js";
 import { t as merge_default } from "./merge-Dl1bfxsj.js";
 import { A as looseObject, B as union, I as record, N as number, P as object, R as string, T as boolean, b as _enum, w as array } from "./zod-DITCj31F.js";

package/dist/{vega-component-2gqZksJH.js → vega-component-rDX7xwxH.js}

@@ -1,19 +1,19 @@
 import { s as __toESM } from "./chunk-BNovOVIE.js";
 import { t as require_react } from "./react-DdA8EBol.js";
 import { t as require_compiler_runtime } from "./compiler-runtime-DHFVbq0b.js";
-import "./Combination-BOmAhdTT.js";
-import { S as CircleQuestionMark, a as AlertTitle, m as asRemoteURL, n as arrow, o as isValid, r as Alert, t as useDeepCompareMemoize } from "./useDeepCompareMemoize-HKwezUBx.js";
-import { l as Events } from "./button-D6ZIdUA3.js";
-import { o as Objects, s as Logger } from "./hotkeys-C4e3s3sJ.js";
+import "./Combination-Bg-xN8JV.js";
+import { S as CircleQuestionMark, a as AlertTitle, m as asRemoteURL, n as arrow, o as isValid, r as Alert, t as useDeepCompareMemoize } from "./useDeepCompareMemoize-CfoxVor3.js";
+import { l as Events } from "./button-BWvsJ2Wr.js";
+import { o as Objects, s as Logger } from "./hotkeys-B5WnGZXF.js";
 import { t as require_jsx_runtime } from "./jsx-runtime-CTBg5pdT.js";
 import "./react-dom-DJW8xUDg.js";
-import { l as Tooltip, n as ErrorBanner } from "./error-banner-uJ4xh94e.js";
-import { h as useEvent_default, n as useTheme } from "./useTheme-D0zT6HZe.js";
+import { l as Tooltip, n as ErrorBanner } from "./error-banner-C7KLpECd.js";
+import { h as useEvent_default, n as useTheme } from "./useTheme-BYG2SH8J.js";
 import { t as _baseUniq_default } from "./_baseUniq-4lqa8rDi.js";
 import { i as debounce_default } from "./constants-DrOu5vvd.js";
-import { a as tooltipHandler, n as vegaLoadData } from "./loader-WBQoIunJ.js";
+import { a as tooltipHandler, n as vegaLoadData } from "./loader-B0KEFFi-.js";
 import { n as formats } from "./vega-loader.browser-BJ1uJidF.js";
-import { t as useAsyncData } from "./useAsyncData-DFBEii2l.js";
+import { t as useAsyncData } from "./useAsyncData-BjNwqCfS.js";
 import { t as j } from "./react-vega-CjiPWyw0.js";
 import "./defaultLocale-qS7DaAmi.js";
 import "./defaultLocale-Bxoo2-30.js";

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@marimo-team/islands",
-  "version": "0.19.7-dev24",
+  "version": "0.19.7-dev26",
   "main": "dist/main.js",
   "types": "dist/index.d.ts",
   "type": "module",

package/src/components/editor/actions/useNotebookActions.tsx

@@ -77,6 +77,7 @@ import {
 } from "@/utils/download";
 import { Filenames } from "@/utils/filenames";
 import { Objects } from "@/utils/objects";
+import type { ProgressState } from "@/utils/progress";
 import { newNotebookURL } from "@/utils/urls";
 import { useRunAllCells } from "../cell/useRunCells";
 import { useChromeActions, useChromeState } from "../chrome/state";
@@ -238,11 +239,12 @@ export function useNotebookActions() {
                 return;
               }
 
-              const downloadPDF = async () => {
-                await updateCellOutputsWithScreenshots(
+              const downloadPDF = async (progress: ProgressState) => {
+                await updateCellOutputsWithScreenshots({
+                  progress,
                   takeScreenshots,
                   updateCellOutputs,
-                );
+                });
                 await downloadAsPDF({
                   filename: filename,
                   webpdf: false,

package/src/components/ui/progress.tsx

@@ -5,10 +5,18 @@ import * as React from "react";
 
 import { cn } from "@/utils/cn";
 
+interface ProgressProps
+  extends React.ComponentPropsWithoutRef<typeof ProgressPrimitive.Root> {
+  /**
+   * When true, shows an indeterminate animated progress bar.
+   */
+  indeterminate?: boolean;
+}
+
 const Progress = React.forwardRef<
-  React.ElementRef<typeof ProgressPrimitive.Root>,
-  React.ComponentPropsWithoutRef<typeof ProgressPrimitive.Root>
->(({ className, value, ...props }, ref) => (
+  React.ComponentRef<typeof ProgressPrimitive.Root>,
+  ProgressProps
+>(({ className, value, indeterminate, ...props }, ref) => (
   <ProgressPrimitive.Root
     ref={ref}
     className={cn(
@@ -18,8 +26,17 @@ const Progress = React.forwardRef<
     {...props}
   >
     <ProgressPrimitive.Indicator
-      className="h-full w-full flex-1 bg-primary transition-all pulse"
-      style={{ transform: `translateX(-${100 - (value || 0)}%)` }}
+      className={cn(
+        "h-full flex-1 bg-primary",
+        indeterminate
+          ? "w-1/3 animate-progress-indeterminate"
+          : "w-full transition-transform duration-300 ease-out",
+      )}
+      style={
+        indeterminate
+          ? undefined
+          : { transform: `translateX(-${100 - (value || 0)}%)` }
+      }
     />
   </ProgressPrimitive.Root>
 ));

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

@@ -8,6 +8,7 @@ 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 { ProgressState } from "@/utils/progress";
 import {
   updateCellOutputsWithScreenshots,
   useEnrichCellOutputs,
@@ -38,6 +39,8 @@ vi.mock("@/core/cells/cells", async () => {
   };
 });
 
+const progress = ProgressState.indeterminate();
+
 import { toPng } from "html-to-image";
 import { toast } from "@/components/ui/use-toast";
 import { cellsRuntimeAtom } from "@/core/cells/cells";
@@ -100,7 +103,7 @@ describe("useEnrichCellOutputs", () => {
     const { result } = renderHook(() => useEnrichCellOutputs(), { wrapper });
 
     const enrichCellOutputs = result.current;
-    const output = await enrichCellOutputs();
+    const output = await enrichCellOutputs(progress);
 
     expect(output).toEqual({});
     expect(document.getElementById).not.toHaveBeenCalled();
@@ -132,7 +135,7 @@ describe("useEnrichCellOutputs", () => {
     const { result } = renderHook(() => useEnrichCellOutputs(), { wrapper });
 
     const enrichCellOutputs = result.current;
-    const output = await enrichCellOutputs();
+    const output = await enrichCellOutputs(progress);
 
     expect(document.getElementById).toHaveBeenCalledWith(
       CellOutputId.create(cellId),
@@ -177,7 +180,7 @@ describe("useEnrichCellOutputs", () => {
 
     // First call - should capture
     let enrichCellOutputs = result.current;
-    let output = await enrichCellOutputs();
+    let output = await enrichCellOutputs(progress);
     expect(output).toEqual({ [cellId]: ["image/png", mockDataUrl] });
     expect(toPng).toHaveBeenCalledTimes(1);
 
@@ -186,7 +189,7 @@ describe("useEnrichCellOutputs", () => {
 
     // Second call with same output - should not capture again
     enrichCellOutputs = result.current;
-    output = await enrichCellOutputs();
+    output = await enrichCellOutputs(progress);
     expect(output).toEqual({}); // Empty because output hasn't changed
     expect(toPng).toHaveBeenCalledTimes(1); // Still only 1 call
   });
@@ -215,7 +218,7 @@ describe("useEnrichCellOutputs", () => {
     const { result } = renderHook(() => useEnrichCellOutputs(), { wrapper });
 
     const enrichCellOutputs = result.current;
-    const output = await enrichCellOutputs();
+    const output = await enrichCellOutputs(progress);
 
     expect(output).toEqual({}); // Failed screenshot should be filtered out
     expect(Logger.error).toHaveBeenCalledWith(
@@ -245,7 +248,7 @@ describe("useEnrichCellOutputs", () => {
     const { result } = renderHook(() => useEnrichCellOutputs(), { wrapper });
 
     const enrichCellOutputs = result.current;
-    const output = await enrichCellOutputs();
+    const output = await enrichCellOutputs(progress);
 
     expect(output).toEqual({});
     expect(Logger.error).toHaveBeenCalledWith(
@@ -294,7 +297,7 @@ describe("useEnrichCellOutputs", () => {
     const { result } = renderHook(() => useEnrichCellOutputs(), { wrapper });
 
     const enrichCellOutputs = result.current;
-    const output = await enrichCellOutputs();
+    const output = await enrichCellOutputs(progress);
 
     expect(output).toEqual({
       [cell1]: ["image/png", mockDataUrl1],
@@ -340,7 +343,7 @@ describe("useEnrichCellOutputs", () => {
     const { result } = renderHook(() => useEnrichCellOutputs(), { wrapper });
 
     const enrichCellOutputs = result.current;
-    const output = await enrichCellOutputs();
+    const output = await enrichCellOutputs(progress);
 
     // Only the successful screenshot should be in the result
     expect(output).toEqual({
@@ -382,13 +385,13 @@ describe("useEnrichCellOutputs", () => {
 
     // First screenshot
     let enrichCellOutputs = result.current;
-    let output = await enrichCellOutputs();
+    let output = await enrichCellOutputs(progress);
     expect(output).toEqual({ [cellId]: ["image/png", mockDataUrl1] });
 
     // Second call - same output, should not be captured
     rerender();
     enrichCellOutputs = result.current;
-    output = await enrichCellOutputs();
+    output = await enrichCellOutputs(progress);
     expect(output).toEqual({});
 
     // Third call - output changed, should be captured
@@ -407,7 +410,7 @@ describe("useEnrichCellOutputs", () => {
 
     rerender();
     enrichCellOutputs = result.current;
-    output = await enrichCellOutputs();
+    output = await enrichCellOutputs(progress);
     expect(output).toEqual({ [cellId]: ["image/png", mockDataUrl2] });
     expect(toPng).toHaveBeenCalledTimes(2);
   });
@@ -447,7 +450,7 @@ describe("useEnrichCellOutputs", () => {
     const { result } = renderHook(() => useEnrichCellOutputs(), { wrapper });
 
     const enrichCellOutputs = result.current;
-    const output = await enrichCellOutputs();
+    const output = await enrichCellOutputs(progress);
 
     // None of these should trigger screenshots
     expect(output).toEqual({});
@@ -472,7 +475,7 @@ describe("useEnrichCellOutputs", () => {
     const { result } = renderHook(() => useEnrichCellOutputs(), { wrapper });
 
     const enrichCellOutputs = result.current;
-    const output = await enrichCellOutputs();
+    const output = await enrichCellOutputs(progress);
 
     expect(output).toEqual({});
     expect(document.getElementById).not.toHaveBeenCalled();
@@ -504,7 +507,7 @@ describe("useEnrichCellOutputs", () => {
     const { result } = renderHook(() => useEnrichCellOutputs(), { wrapper });
 
     const enrichCellOutputs = result.current;
-    const output = await enrichCellOutputs();
+    const output = await enrichCellOutputs(progress);
 
     // Verify the exact return type structure
     expect(output).toHaveProperty(cellId);
@@ -535,7 +538,11 @@ describe("updateCellOutputsWithScreenshots", () => {
     const takeScreenshots = vi.fn().mockResolvedValue(mockScreenshots);
     const updateCellOutputs = vi.fn().mockResolvedValue(null);
 
-    await updateCellOutputsWithScreenshots(takeScreenshots, updateCellOutputs);
+    await updateCellOutputsWithScreenshots({
+      progress,
+      takeScreenshots,
+      updateCellOutputs,
+    });
 
     expect(takeScreenshots).toHaveBeenCalledTimes(1);
     expect(updateCellOutputs).toHaveBeenCalledTimes(1);
@@ -548,7 +555,11 @@ describe("updateCellOutputsWithScreenshots", () => {
     const takeScreenshots = vi.fn().mockResolvedValue({});
     const updateCellOutputs = vi.fn().mockResolvedValue(null);
 
-    await updateCellOutputsWithScreenshots(takeScreenshots, updateCellOutputs);
+    await updateCellOutputsWithScreenshots({
+      progress,
+      takeScreenshots,
+      updateCellOutputs,
+    });
 
     expect(takeScreenshots).toHaveBeenCalledTimes(1);
     expect(updateCellOutputs).not.toHaveBeenCalled();
@@ -571,7 +582,11 @@ describe("updateCellOutputsWithScreenshots", () => {
     const takeScreenshots = vi.fn().mockResolvedValue(mockScreenshots);
     const updateCellOutputs = vi.fn().mockResolvedValue(null);
 
-    await updateCellOutputsWithScreenshots(takeScreenshots, updateCellOutputs);
+    await updateCellOutputsWithScreenshots({
+      progress,
+      takeScreenshots,
+      updateCellOutputs,
+    });
 
     expect(updateCellOutputs).toHaveBeenCalledWith({
       cellIdsToOutput: mockScreenshots,
@@ -584,7 +599,11 @@ describe("updateCellOutputsWithScreenshots", () => {
     const updateCellOutputs = vi.fn().mockResolvedValue(null);
 
     // Should not throw - errors are caught and shown via toast
-    await updateCellOutputsWithScreenshots(takeScreenshots, updateCellOutputs);
+    await updateCellOutputsWithScreenshots({
+      progress,
+      takeScreenshots,
+      updateCellOutputs,
+    });
 
     expect(updateCellOutputs).not.toHaveBeenCalled();
     expect(Logger.error).toHaveBeenCalledWith(
@@ -613,7 +632,11 @@ describe("updateCellOutputsWithScreenshots", () => {
     const updateCellOutputs = vi.fn().mockRejectedValue(error);
 
     // Should not throw - errors are caught and shown via toast
-    await updateCellOutputsWithScreenshots(takeScreenshots, updateCellOutputs);
+    await updateCellOutputsWithScreenshots({
+      progress,
+      takeScreenshots,
+      updateCellOutputs,
+    });
 
     expect(Logger.error).toHaveBeenCalledWith(
       "Error updating cell outputs with screenshots:",

package/src/core/export/hooks.ts

@@ -7,6 +7,7 @@ import { useInterval } from "@/hooks/useInterval";
 import { getImageDataUrlForCell } from "@/utils/download";
 import { Logger } from "@/utils/Logger";
 import { Objects } from "@/utils/objects";
+import { ProgressState } from "@/utils/progress";
 import { cellsRuntimeAtom } from "../cells/cells";
 import type { CellId } from "../cells/ids";
 import { connectionAtom } from "../network/connection";
@@ -63,10 +64,11 @@ export function useAutoExport() {
 
   useInterval(
     async () => {
-      await updateCellOutputsWithScreenshots(
+      await updateCellOutputsWithScreenshots({
+        progress: ProgressState.indeterminate(),
         takeScreenshots,
         updateCellOutputs,
-      );
+      });
       await autoExportAsIPYNB({
         download: false,
       });
@@ -95,6 +97,8 @@ const MIME_TYPES_TO_CAPTURE_SCREENSHOTS = new Set<MimeType>([
   "application/vnd.vega.v6+json",
 ]);
 
+type ScreenshotResults = Record<CellId, ["image/png", string]>;
+
 /**
  * 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.
@@ -103,7 +107,7 @@ export function useEnrichCellOutputs() {
   const [richCellsOutput, setRichCellsOutput] = useAtom(richCellsToOutputAtom);
   const cellRuntimes = useAtomValue(cellsRuntimeAtom);
 
-  return async (): Promise<Record<CellId, ["image/png", string]>> => {
+  return async (progress: ProgressState): Promise<ScreenshotResults> => {
     const trackedCellsOutput: Record<CellId, unknown> = {};
 
     const cellsToCaptureScreenshot: [CellId, unknown][] = [];
@@ -129,43 +133,40 @@ export function useEnrichCellOutputs() {
     }
 
     // Capture screenshots
-    const results = await Promise.all(
-      cellsToCaptureScreenshot.map(async ([cellId]) => {
-        try {
-          const dataUrl = await getImageDataUrlForCell(cellId, false);
-          if (!dataUrl) {
-            Logger.error(`Failed to capture screenshot for cell ${cellId}`);
-            return null;
-          }
-          return [cellId, ["image/png", dataUrl]] as [
-            CellId,
-            ["image/png", string],
-          ];
-        } catch (error) {
-          Logger.error(`Error screenshotting cell ${cellId}:`, error);
-          return null;
+    const total = cellsToCaptureScreenshot.length;
+    progress.addTotal(total);
+    const results: ScreenshotResults = {};
+    for (const [cellId] of cellsToCaptureScreenshot) {
+      try {
+        const dataUrl = await getImageDataUrlForCell(cellId, false);
+        if (!dataUrl) {
+          Logger.error(`Failed to capture screenshot for cell ${cellId}`);
+          continue;
         }
-      }),
-    );
-
-    return Objects.fromEntries(
-      results.filter(
-        (result): result is [CellId, ["image/png", string]] => result !== null,
-      ),
-    );
+        results[cellId] = ["image/png", dataUrl];
+      } catch (error) {
+        Logger.error(`Error screenshotting cell ${cellId}:`, error);
+      } finally {
+        progress.increment(1);
+      }
+    }
+
+    return results;
   };
 }
 
 /**
  * Utility function to take screenshots of cells with HTML outputs and update the cell outputs.
  */
-export async function updateCellOutputsWithScreenshots(
-  takeScreenshots: () => Promise<Record<CellId, ["image/png", string]>>,
-  updateCellOutputs: (request: UpdateCellOutputsRequest) => Promise<null>,
-) {
+export async function updateCellOutputsWithScreenshots(opts: {
+  progress: ProgressState;
+  takeScreenshots: (progress: ProgressState) => Promise<ScreenshotResults>;
+  updateCellOutputs: (request: UpdateCellOutputsRequest) => Promise<null>;
+}) {
+  const { progress, takeScreenshots, updateCellOutputs } = opts;
   try {
-    const cellIdsToOutput = await takeScreenshots();
-    if (Object.keys(cellIdsToOutput).length > 0) {
+    const cellIdsToOutput = await takeScreenshots(progress);
+    if (Objects.size(cellIdsToOutput) > 0) {
       await updateCellOutputs({ cellIdsToOutput });
     }
   } catch (error) {

package/src/utils/__tests__/download.test.tsx

@@ -57,10 +57,12 @@ describe("withLoadingToast", () => {
     });
 
     expect(toast).toHaveBeenCalledTimes(1);
-    expect(toast).toHaveBeenCalledWith({
-      title: "Loading...",
-      duration: Infinity,
-    });
+    expect(toast).toHaveBeenCalledWith(
+      expect.objectContaining({
+        title: "Loading...",
+        duration: Infinity,
+      }),
+    );
     expect(mockDismiss).toHaveBeenCalledTimes(1);
     expect(result).toBe("success");
   });