@marimo-team/frontend 0.19.7-dev35 → 0.19.7-dev37

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/core/network/types.ts

@@ -100,6 +100,10 @@ export type InvokeAiToolRequest = schemas["InvokeAiToolRequest"];
 export type InvokeAiToolResponse = schemas["InvokeAiToolResponse"];
 export type ClearCacheRequest = schemas["ClearCacheRequest"];
 export type GetCacheInfoRequest = schemas["GetCacheInfoRequest"];
+export type LspHealthResponse = schemas["LspHealthResponse"];
+export type LspRestartRequest = schemas["LspRestartRequest"];
+export type LspRestartResponse = schemas["LspRestartResponse"];
+export type LspServerHealth = schemas["LspServerHealth"];
 
 /**
  * Requests sent to the BE during run/edit mode.

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;

package/src/utils/__tests__/async-capture-tracker.test.ts

@@ -0,0 +1,353 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+import { beforeEach, describe, expect, it } from "vitest";
+import { AsyncCaptureTracker } from "../async-capture-tracker";
+
+describe("AsyncCaptureTracker", () => {
+  let tracker: AsyncCaptureTracker<string, string>;
+
+  beforeEach(() => {
+    tracker = new AsyncCaptureTracker();
+  });
+
+  describe("needsCapture", () => {
+    it("should return true for keys not yet captured", () => {
+      expect(tracker.needsCapture("a", "value1")).toBe(true);
+    });
+
+    it("should return false for keys captured with the same value", () => {
+      const handle = tracker.startCapture("a", "value1");
+      handle.markCaptured("result");
+      expect(tracker.needsCapture("a", "value1")).toBe(false);
+    });
+
+    it("should return true for keys captured with a different value", () => {
+      const handle = tracker.startCapture("a", "value1");
+      handle.markCaptured("result");
+      expect(tracker.needsCapture("a", "value2")).toBe(true);
+    });
+
+    it("should return false for keys in-flight with the same value", () => {
+      tracker.startCapture("a", "value1");
+      expect(tracker.needsCapture("a", "value1")).toBe(false);
+    });
+
+    it("should return true for keys in-flight with a different value", () => {
+      tracker.startCapture("a", "value1");
+      expect(tracker.needsCapture("a", "value2")).toBe(true);
+    });
+
+    it("should return true for keys that failed (returned to idle)", () => {
+      const handle = tracker.startCapture("a", "value1");
+      handle.markFailed();
+      expect(tracker.needsCapture("a", "value1")).toBe(true);
+    });
+  });
+
+  describe("startCapture", () => {
+    it("should return a handle with a non-aborted signal", () => {
+      const handle = tracker.startCapture("a", "v");
+      expect(handle.signal.aborted).toBe(false);
+    });
+
+    it("should abort only the same key's previous signal", () => {
+      const handleA = tracker.startCapture("a", "v1");
+      const handleB = tracker.startCapture("b", "v1");
+      const handleA2 = tracker.startCapture("a", "v2");
+
+      expect(handleA.signal.aborted).toBe(true);
+      expect(handleB.signal.aborted).toBe(false);
+      expect(handleA2.signal.aborted).toBe(false);
+
+      // Unrelated key still completes independently
+      handleB.markCaptured("result-b");
+      expect(tracker.needsCapture("b", "v1")).toBe(false);
+    });
+
+    it("should not abort other keys when one is re-captured", () => {
+      tracker.startCapture("a", "v");
+      tracker.startCapture("b", "v");
+      tracker.startCapture("c", "v");
+
+      const handleB2 = tracker.startCapture("b", "v2");
+
+      expect(tracker.needsCapture("a", "v")).toBe(false); // still in-flight
+      expect(tracker.needsCapture("c", "v")).toBe(false); // still in-flight
+      expect(handleB2.signal.aborted).toBe(false);
+    });
+
+    it("should mark key as in-flight", () => {
+      tracker.startCapture("a", "v");
+      expect(tracker.isCapturing).toBe(true);
+      expect(tracker.needsCapture("a", "v")).toBe(false);
+    });
+  });
+
+  describe("handle.markCaptured", () => {
+    it("should remove key from in-flight", () => {
+      const handle = tracker.startCapture("a", "v");
+      expect(tracker.isCapturing).toBe(true);
+      handle.markCaptured("result");
+      expect(tracker.isCapturing).toBe(false);
+    });
+
+    it("should persist the captured value", () => {
+      const handle = tracker.startCapture("a", "value1");
+      handle.markCaptured("result");
+      expect(tracker.needsCapture("a", "value1")).toBe(false);
+      expect(tracker.needsCapture("a", "value2")).toBe(true);
+    });
+
+    it("should update the captured value on re-capture", () => {
+      const h1 = tracker.startCapture("a", "v1");
+      h1.markCaptured("r1");
+
+      const h2 = tracker.startCapture("a", "v2");
+      h2.markCaptured("r2");
+
+      expect(tracker.needsCapture("a", "v1")).toBe(true);
+      expect(tracker.needsCapture("a", "v2")).toBe(false);
+    });
+
+    it("should be a no-op when handle is stale (superseded)", () => {
+      const handle1 = tracker.startCapture("a", "v1");
+      const handle2 = tracker.startCapture("a", "v2");
+
+      // handle1 is stale — markCaptured should not affect state
+      handle1.markCaptured("stale-result");
+
+      // "a" should still be in-flight with v2
+      expect(tracker.needsCapture("a", "v2")).toBe(false); // in-flight
+      expect(tracker.needsCapture("a", "v1")).toBe(true); // not captured
+      expect(tracker.isCapturing).toBe(true);
+
+      // handle2 still works
+      handle2.markCaptured("real-result");
+      expect(tracker.needsCapture("a", "v2")).toBe(false); // captured
+      expect(tracker.isCapturing).toBe(false);
+    });
+  });
+
+  describe("handle.markFailed", () => {
+    it("should remove key from in-flight", () => {
+      const handle = tracker.startCapture("a", "v");
+      handle.markFailed();
+      expect(tracker.isCapturing).toBe(false);
+    });
+
+    it("should not add a captured entry", () => {
+      const handle = tracker.startCapture("a", "v");
+      handle.markFailed();
+      expect(tracker.needsCapture("a", "v")).toBe(true);
+    });
+
+    it("should not clear a previously captured value", () => {
+      const h1 = tracker.startCapture("a", "v1");
+      h1.markCaptured("r1");
+
+      const h2 = tracker.startCapture("a", "v2");
+      h2.markFailed();
+
+      expect(tracker.needsCapture("a", "v1")).toBe(false);
+      expect(tracker.needsCapture("a", "v2")).toBe(true);
+    });
+
+    it("should be a no-op when handle is stale", () => {
+      const handle1 = tracker.startCapture("a", "v1");
+      tracker.startCapture("a", "v2");
+
+      handle1.markFailed();
+
+      // "a" should still be in-flight with v2
+      expect(tracker.isCapturing).toBe(true);
+      expect(tracker.needsCapture("a", "v2")).toBe(false);
+    });
+  });
+
+  describe("waitForInFlight", () => {
+    it("should return a promise for in-flight key with same value", async () => {
+      const handle = tracker.startCapture("a", "v1");
+      const promise = tracker.waitForInFlight("a", "v1");
+      expect(promise).not.toBeNull();
+
+      handle.markCaptured("result-a");
+      const result = await promise;
+      expect(result).toBe("result-a");
+    });
+
+    it("should return null for captured key (not in-flight)", () => {
+      const handle = tracker.startCapture("a", "v1");
+      handle.markCaptured("result");
+      expect(tracker.waitForInFlight("a", "v1")).toBeNull();
+    });
+
+    it("should return null for key not being tracked", () => {
+      expect(tracker.waitForInFlight("a", "v1")).toBeNull();
+    });
+
+    it("should return null for in-flight key with different value", () => {
+      tracker.startCapture("a", "v1");
+      expect(tracker.waitForInFlight("a", "v2")).toBeNull();
+    });
+
+    it("should resolve with undefined when capture fails", async () => {
+      const handle = tracker.startCapture("a", "v1");
+      const promise = tracker.waitForInFlight("a", "v1");
+
+      handle.markFailed();
+      const result = await promise;
+      expect(result).toBeUndefined();
+    });
+
+    it("should resolve with undefined when capture is superseded", async () => {
+      tracker.startCapture("a", "v1");
+      const promise = tracker.waitForInFlight("a", "v1");
+
+      // Supersede with new capture
+      tracker.startCapture("a", "v2");
+      const result = await promise;
+      expect(result).toBeUndefined();
+    });
+
+    it("should resolve with undefined when aborted", async () => {
+      tracker.startCapture("a", "v1");
+      const promise = tracker.waitForInFlight("a", "v1");
+
+      tracker.abort();
+      const result = await promise;
+      expect(result).toBeUndefined();
+    });
+
+    it("should allow multiple waiters on the same key", async () => {
+      const handle = tracker.startCapture("a", "v1");
+      const p1 = tracker.waitForInFlight("a", "v1");
+      const p2 = tracker.waitForInFlight("a", "v1");
+
+      handle.markCaptured("shared-result");
+
+      const [r1, r2] = await Promise.all([p1, p2]);
+      expect(r1).toBe("shared-result");
+      expect(r2).toBe("shared-result");
+    });
+  });
+
+  describe("prune", () => {
+    it("should remove captured entries for keys not in the set", () => {
+      const hA = tracker.startCapture("a", "v1");
+      hA.markCaptured("r1");
+      const hB = tracker.startCapture("b", "v2");
+      hB.markCaptured("r2");
+
+      tracker.prune(new Set(["a"]));
+
+      expect(tracker.needsCapture("a", "v1")).toBe(false);
+      expect(tracker.needsCapture("b", "v2")).toBe(true);
+    });
+
+    it("should abort and remove in-flight entries for pruned keys", () => {
+      const handleA = tracker.startCapture("a", "v");
+      const handleB = tracker.startCapture("b", "v");
+
+      tracker.prune(new Set(["a"]));
+
+      expect(handleA.signal.aborted).toBe(false);
+      expect(handleB.signal.aborted).toBe(true);
+      expect(tracker.needsCapture("b", "v")).toBe(true);
+    });
+
+    it("should resolve pruned in-flight deferreds with undefined", async () => {
+      tracker.startCapture("b", "v");
+      const promise = tracker.waitForInFlight("b", "v");
+
+      tracker.prune(new Set(["a"]));
+      const result = await promise;
+      expect(result).toBeUndefined();
+    });
+
+    it("should do nothing if all keys are current", () => {
+      const hA = tracker.startCapture("a", "v1");
+      hA.markCaptured("r1");
+      const hB = tracker.startCapture("b", "v2");
+      hB.markCaptured("r2");
+
+      tracker.prune(new Set(["a", "b"]));
+
+      expect(tracker.needsCapture("a", "v1")).toBe(false);
+      expect(tracker.needsCapture("b", "v2")).toBe(false);
+    });
+  });
+
+  describe("abort", () => {
+    it("should abort all in-flight signals", () => {
+      const handleA = tracker.startCapture("a", "v");
+      const handleB = tracker.startCapture("b", "v");
+      tracker.abort();
+      expect(handleA.signal.aborted).toBe(true);
+      expect(handleB.signal.aborted).toBe(true);
+    });
+
+    it("should clear in-flight state", () => {
+      tracker.startCapture("a", "v");
+      tracker.abort();
+      expect(tracker.isCapturing).toBe(false);
+      expect(tracker.needsCapture("a", "v")).toBe(true);
+    });
+
+    it("should not clear captured state", () => {
+      const handle = tracker.startCapture("a", "v1");
+      handle.markCaptured("result");
+      tracker.abort();
+      expect(tracker.needsCapture("a", "v1")).toBe(false);
+    });
+  });
+
+  describe("reset", () => {
+    it("should clear all state", () => {
+      const hA = tracker.startCapture("a", "v1");
+      hA.markCaptured("r1");
+      tracker.startCapture("b", "v2");
+
+      tracker.reset();
+
+      expect(tracker.isCapturing).toBe(false);
+      expect(tracker.needsCapture("a", "v1")).toBe(true);
+      expect(tracker.needsCapture("b", "v2")).toBe(true);
+    });
+  });
+
+  describe("isCapturing", () => {
+    it("should be false when nothing is in-flight", () => {
+      expect(tracker.isCapturing).toBe(false);
+    });
+
+    it("should be true when items are in-flight", () => {
+      tracker.startCapture("a", "v");
+      expect(tracker.isCapturing).toBe(true);
+    });
+
+    it("should be false when all items are completed", () => {
+      const hA = tracker.startCapture("a", "v");
+      const hB = tracker.startCapture("b", "v");
+      hA.markCaptured("r");
+      hB.markFailed();
+      expect(tracker.isCapturing).toBe(false);
+    });
+  });
+
+  describe("granular abort", () => {
+    it("should allow independent completion of unrelated keys", () => {
+      const hA = tracker.startCapture("a", "v1");
+      tracker.startCapture("b", "v1");
+      const hC = tracker.startCapture("c", "v1");
+
+      hA.markCaptured("ra");
+      tracker.startCapture("b", "v2"); // abort old "b", start new
+      hC.markCaptured("rc");
+
+      expect(tracker.needsCapture("a", "v1")).toBe(false);
+      expect(tracker.needsCapture("c", "v1")).toBe(false);
+      expect(tracker.needsCapture("b", "v2")).toBe(false); // in-flight
+      expect(tracker.needsCapture("b", "v1")).toBe(true);
+    });
+  });
+});