Ensure overflow is captured in screenshots, and ensure smooth capture (#8046)

Author: Light2Dark

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

@@ -166,13 +166,63 @@ describe("getImageDataUrlForCell", () => {
     );
   });
 
-  it("should restore original overflow style after capture", async () => {
+  it("should pass style options to prevent clipping", async () => {
+    vi.mocked(toPng).mockResolvedValue(mockDataUrl);
+
+    await getImageDataUrlForCell("cell-1" as CellId);
+
+    expect(toPng).toHaveBeenCalledWith(
+      mockElement,
+      expect.objectContaining({
+        style: {
+          maxHeight: "none",
+          overflow: "visible",
+        },
+      }),
+    );
+  });
+
+  it("should pass scrollHeight as height option", async () => {
+    // Set up element with scrollHeight
+    Object.defineProperty(mockElement, "scrollHeight", {
+      value: 500,
+      configurable: true,
+    });
+    vi.mocked(toPng).mockResolvedValue(mockDataUrl);
+
+    await getImageDataUrlForCell("cell-1" as CellId);
+
+    expect(toPng).toHaveBeenCalledWith(
+      mockElement,
+      expect.objectContaining({
+        height: 500,
+      }),
+    );
+  });
+
+  it("should pass scrollbar hiding styles via extraStyleContent", async () => {
+    vi.mocked(toPng).mockResolvedValue(mockDataUrl);
+
+    await getImageDataUrlForCell("cell-1" as CellId);
+
+    expect(toPng).toHaveBeenCalledWith(
+      mockElement,
+      expect.objectContaining({
+        extraStyleContent: expect.stringContaining("scrollbar-width: none"),
+      }),
+    );
+  });
+
+  it("should not modify the live DOM element", async () => {
     mockElement.style.overflow = "hidden";
+    mockElement.style.maxHeight = "100px";
     vi.mocked(toPng).mockResolvedValue(mockDataUrl);
 
     await getImageDataUrlForCell("cell-1" as CellId);
 
+    // DOM should remain unchanged
     expect(mockElement.style.overflow).toBe("hidden");
+    expect(mockElement.style.maxHeight).toBe("100px");
   });
 
   it("should throw error on failure", async () => {
@@ -183,34 +233,20 @@ describe("getImageDataUrlForCell", () => {
     );
   });
 
-  it("should cleanup even on failure", async () => {
-    mockElement.style.overflow = "scroll";
-    vi.mocked(toPng).mockRejectedValue(new Error("Capture failed"));
-
-    await expect(getImageDataUrlForCell("cell-1" as CellId)).rejects.toThrow();
-
-    expect(document.body.classList.contains("printing")).toBe(false);
-    expect(mockElement.style.overflow).toBe("scroll");
-  });
-
   it("should handle concurrent captures correctly", async () => {
     // Create a second element
     const mockElement2 = document.createElement("div");
     mockElement2.id = CellOutputId.create("cell-2" as CellId);
     document.body.append(mockElement2);
 
-    vi.mocked(toPng).mockImplementation(async () => {
-      // body.printing should not be added during cell captures
-      expect(document.body.classList.contains("printing")).toBe(false);
-      return mockDataUrl;
-    });
+    vi.mocked(toPng).mockResolvedValue(mockDataUrl);
 
     const capture1 = getImageDataUrlForCell("cell-1" as CellId);
     const capture2 = getImageDataUrlForCell("cell-2" as CellId);
 
     await Promise.all([capture1, capture2]);
 
-    expect(document.body.classList.contains("printing")).toBe(false);
+    expect(toPng).toHaveBeenCalledTimes(2);
 
     mockElement2.remove();
   });
@@ -266,23 +302,6 @@ describe("downloadHTMLAsImage", () => {
     expect(mockAnchor.click).toHaveBeenCalled();
   });
 
-  it("should add body.printing class without prepare function", async () => {
-    vi.mocked(toPng).mockImplementation(async () => {
-      expect(document.body.classList.contains("printing")).toBe(true);
-      return mockDataUrl;
-    });
-
-    await downloadHTMLAsImage({ element: mockElement, filename: "test" });
-  });
-
-  it("should remove body.printing class after download without prepare", async () => {
-    vi.mocked(toPng).mockResolvedValue(mockDataUrl);
-
-    await downloadHTMLAsImage({ element: mockElement, filename: "test" });
-
-    expect(document.body.classList.contains("printing")).toBe(false);
-  });
-
   it("should use prepare function when provided", async () => {
     vi.mocked(toPng).mockResolvedValue(mockDataUrl);
     const cleanup = vi.fn();
@@ -406,24 +425,32 @@ describe("downloadCellOutputAsImage", () => {
     expect(mockAnchor.download).toBe("result.png");
   });
 
-  it("should apply cell-specific preparation", async () => {
-    vi.mocked(toPng).mockImplementation(async () => {
-      // Check that cell-specific classes are applied
-      expect(mockElement.style.overflow).toBe("visible");
-      return mockDataUrl;
-    });
+  it("should pass style options to toPng for full content capture", async () => {
+    vi.mocked(toPng).mockResolvedValue(mockDataUrl);
 
     await downloadCellOutputAsImage("cell-1" as CellId, "result");
+
+    expect(toPng).toHaveBeenCalledWith(
+      mockElement,
+      expect.objectContaining({
+        style: {
+          maxHeight: "none",
+          overflow: "visible",
+        },
+      }),
+    );
   });
 
-  it("should cleanup after download", async () => {
-    mockElement.style.overflow = "visible";
+  it("should not modify the live DOM element", async () => {
+    mockElement.style.overflow = "hidden";
+    mockElement.style.maxHeight = "100px";
     vi.mocked(toPng).mockResolvedValue(mockDataUrl);
 
     await downloadCellOutputAsImage("cell-1" as CellId, "result");
 
-    expect(document.body.classList.contains("printing")).toBe(false);
-    expect(mockElement.style.overflow).toBe("visible");
+    // DOM should remain unchanged
+    expect(mockElement.style.overflow).toBe("hidden");
+    expect(mockElement.style.maxHeight).toBe("100px");
   });
 });
 

frontend/src/utils/download.ts

@@ -46,24 +46,6 @@ function findElementForCell(cellId: CellId): HTMLElement | undefined {
   return element;
 }
 
-/**
- * Prepare a cell element for screenshot capture.
- *
- * @param element - The cell output element to prepare
- * @returns A cleanup function to restore the element's original state
- */
-function prepareCellElementForScreenshot(element: HTMLElement) {
-  const originalOverflow = element.style.overflow;
-  const maxHeight = element.style.maxHeight;
-  element.style.overflow = "visible";
-  element.style.maxHeight = "none";
-
-  return () => {
-    element.style.overflow = originalOverflow;
-    element.style.maxHeight = maxHeight;
-  };
-}
-
 const THRESHOLD_TIME_MS = 500;
 const HIDE_SCROLLBAR_STYLES = `
   * { scrollbar-width: none; -ms-overflow-style: none; }
@@ -89,25 +71,26 @@ export async function getImageDataUrlForCell(
     return iframeDataUrl;
   }
 
-  const cleanup = prepareCellElementForScreenshot(element);
-
-  try {
-    const startTime = Date.now();
-    const dataUrl = await toPng(element, {
-      extraStyleContent: HIDE_SCROLLBAR_STYLES,
-    });
-    const timeTaken = Date.now() - startTime;
-    if (timeTaken > THRESHOLD_TIME_MS) {
-      Logger.debug(
-        "toPng operation for element",
-        element,
-        `took ${timeTaken} ms (exceeds threshold)`,
-      );
-    }
-    return dataUrl;
-  } finally {
-    cleanup();
+  const startTime = Date.now();
+  const dataUrl = await toPng(element, {
+    extraStyleContent: HIDE_SCROLLBAR_STYLES,
+    // Add these styles so the element output is not clipped
+    // Width can be clipped since pdf has limited width
+    style: {
+      maxHeight: "none",
+      overflow: "visible",
+    },
+    height: element.scrollHeight,
+  });
+  const timeTaken = Date.now() - startTime;
+  if (timeTaken > THRESHOLD_TIME_MS) {
+    Logger.debug(
+      "toPng operation for element",
+      element,
+      `took ${timeTaken} ms (exceeds threshold)`,
+    );
   }
+  return dataUrl;
 }
 
 /**