@marimo-team/islands 0.19.5-dev40 → 0.19.5-dev44

package/src/components/chat/chat-panel.tsx

@@ -426,14 +426,14 @@ const ChatInput: React.FC<ChatInputProps> = memo(
 ChatInput.displayName = "ChatInput";
 
 const ChatPanel = () => {
-  const aiEnabled = useAtomValue(aiEnabledAtom);
+  const aiConfigured = useAtomValue(aiEnabledAtom);
   const { handleClick } = useOpenSettingsToTab();
 
-  if (!aiEnabled) {
+  if (!aiConfigured) {
     return (
       <PanelEmptyState
         title="Chat with AI"
-        description="AI is currently disabled. Add your API key to enable."
+        description="No AI provider configured or model selected"
         action={
           <Button variant="outline" size="sm" onClick={() => handleClick("ai")}>
             Edit AI settings

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

@@ -26,7 +26,6 @@ import {
   ZapIcon,
   ZapOffIcon,
 } from "lucide-react";
-import { downloadCellOutput } from "@/components/export/export-output-button";
 import { MultiIcon } from "@/components/icons/multi-icon";
 import { useImperativeModal } from "@/components/modal/ImperativeModal";
 import {
@@ -54,6 +53,7 @@ import { useRequestClient } from "@/core/network/requests";
 import type { CellConfig, RuntimeState } from "@/core/network/types";
 import { canLinkToCell, createCellLink } from "@/utils/cell-urls";
 import { copyToClipboard } from "@/utils/copy";
+import { downloadCellOutputAsImage } from "@/utils/download";
 import { MarkdownIcon, PythonIcon } from "../cell/code/icons";
 import { useDeleteCellCallback } from "../cell/useDeleteCell";
 import { useRunCell } from "../cell/useRunCells";
@@ -341,7 +341,7 @@ export function useCellActionButtons({ cell, closePopover }: Props) {
         icon: <ImageIcon size={13} strokeWidth={1.5} />,
         label: "Export output as PNG",
         hidden: !hasOutput,
-        handle: () => downloadCellOutput(cellId),
+        handle: () => downloadCellOutputAsImage(cellId, "result"),
       },
       {
         icon: <XCircleIcon size={13} strokeWidth={1.5} />,

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

@@ -134,6 +134,12 @@ export function useNotebookActions() {
   const sharingHtmlEnabled = resolvedConfig.sharing?.html ?? true;
   const sharingWasmEnabled = resolvedConfig.sharing?.wasm ?? true;
 
+  const isServerSidePdfExportEnabled = getFeatureFlag("server_side_pdf_export");
+  // With server side pdf export, it doesn't matter what mode we are in,
+  // Default export uses browser print, which is better in present mode
+  const pdfDownloadEnabled =
+    isServerSidePdfExportEnabled || viewState.mode === "present";
+
   const renderCheckboxElement = (checked: boolean) => (
     <div className="w-8 flex justify-end">
       {checked && <CheckIcon size={14} />}
@@ -209,22 +215,24 @@ export function useNotebookActions() {
             if (!app) {
               return;
             }
-            await downloadHTMLAsImage(app, document.title);
+            await downloadHTMLAsImage({
+              element: app,
+              filename: document.title,
+            });
           },
         },
         {
           icon: <FileIcon size={14} strokeWidth={1.5} />,
           label: "Download as PDF",
-          disabled: viewState.mode !== "present",
-          tooltip:
-            viewState.mode === "present" ? undefined : (
-              <span>
-                Only available in app view. <br />
-                Toggle with: {renderShortcut("global.hideCode", false)}
-              </span>
-            ),
+          disabled: !pdfDownloadEnabled,
+          tooltip: pdfDownloadEnabled ? undefined : (
+            <span>
+              Only available in app view. <br />
+              Toggle with: {renderShortcut("global.hideCode", false)}
+            </span>
+          ),
           handle: async () => {
-            if (getFeatureFlag("server_side_pdf_export")) {
+            if (isServerSidePdfExportEnabled) {
               if (!filename) {
                 toastNotebookMustBeNamed();
                 return;

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

@@ -185,7 +185,7 @@ const ActionButtons: React.FC<{
     if (!app) {
       return;
     }
-    await downloadHTMLAsImage(app, document.title);
+    await downloadHTMLAsImage({ element: app, filename: document.title });
   };
 
   const handleDownloadAsHTML = async () => {

package/src/core/ai/model-registry.ts

@@ -21,8 +21,17 @@ export interface AiModel extends AiModelType {
   custom: boolean;
 }
 
-const getKnownModelMap = once((): ReadonlyMap<QualifiedModelId, AiModel> => {
+interface KnownModelMaps {
+  /** Map of qualified model ID to model info */
+  modelMap: ReadonlyMap<QualifiedModelId, AiModel>;
+  /** Map of provider ID to first default model (supports chat or edit) */
+  defaultModelByProvider: ReadonlyMap<ProviderId, QualifiedModelId>;
+}
+
+export const getKnownModelMaps = once((): KnownModelMaps => {
   const modelMap = new Map<QualifiedModelId, AiModel>();
+  const defaultModelByProvider = new Map<ProviderId, QualifiedModelId>();
+
   for (const model of models) {
     const modelId = model.model as ShortModelId;
     const modelInfo: AiModel = {
@@ -33,12 +42,21 @@ const getKnownModelMap = once((): ReadonlyMap<QualifiedModelId, AiModel> => {
       custom: false,
     };
 
+    const supportsChatOrEdit =
+      modelInfo.roles.includes("chat") || modelInfo.roles.includes("edit");
+
     for (const provider of modelInfo.providers) {
       const qualifiedModelId: QualifiedModelId = `${provider}/${modelId}`;
       modelMap.set(qualifiedModelId, modelInfo);
+
+      // Track first model per provider that supports chat or edit
+      if (supportsChatOrEdit && !defaultModelByProvider.has(provider)) {
+        defaultModelByProvider.set(provider, qualifiedModelId);
+      }
     }
   }
-  return modelMap;
+
+  return { modelMap, defaultModelByProvider };
 });
 
 const getProviderMap = once(
@@ -125,7 +143,7 @@ export class AiModelRegistry {
   }) {
     const { displayedModels, customModels } = opts;
     const hasDisplayedModels = displayedModels.size > 0;
-    const knownModelMap = getKnownModelMap();
+    const knownModelMap = getKnownModelMaps().modelMap;
     const customModelsMap = new Map<QualifiedModelId, AiModel>();
 
     let modelsMap = new Map<QualifiedModelId, AiModel>();

package/src/core/export/hooks.ts

@@ -1,14 +1,14 @@
 /* Copyright 2026 Marimo. All rights reserved. */
-import { toPng } from "html-to-image";
 import { atom, useAtom, 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 { getImageDataUrlForCell } from "@/utils/download";
 import { Logger } from "@/utils/Logger";
 import { Objects } from "@/utils/objects";
 import { cellsRuntimeAtom } from "../cells/cells";
-import { type CellId, CellOutputId } from "../cells/ids";
+import type { CellId } from "../cells/ids";
 import { connectionAtom } from "../network/connection";
 import { useRequestClient } from "../network/requests";
 import type { UpdateCellOutputsRequest } from "../network/types";
@@ -131,16 +131,12 @@ export function useEnrichCellOutputs() {
     // 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);
+          const dataUrl = await getImageDataUrlForCell(cellId);
+          if (!dataUrl) {
+            Logger.error(`Failed to capture screenshot for cell ${cellId}`);
+            return null;
+          }
           return [cellId, ["image/png", dataUrl]] as [
             CellId,
             ["image/png", string],

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

@@ -1,6 +1,19 @@
 /* Copyright 2026 Marimo. All rights reserved. */
-import { beforeEach, describe, expect, it, vi } from "vitest";
-import { withLoadingToast } from "../download";
+import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
+import type { CellId } from "@/core/cells/ids";
+import { CellOutputId } from "@/core/cells/ids";
+import {
+  downloadByURL,
+  downloadCellOutputAsImage,
+  downloadHTMLAsImage,
+  getImageDataUrlForCell,
+  withLoadingToast,
+} from "../download";
+
+// Mock html-to-image
+vi.mock("html-to-image", () => ({
+  toPng: vi.fn(),
+}));
 
 // Mock the toast module
 const mockDismiss = vi.fn();
@@ -15,7 +28,23 @@ vi.mock("@/components/icons/spinner", () => ({
   Spinner: () => "MockSpinner",
 }));
 
+// Mock Logger
+vi.mock("@/utils/Logger", () => ({
+  Logger: {
+    error: vi.fn(),
+  },
+}));
+
+// Mock Filenames
+vi.mock("@/utils/filenames", () => ({
+  Filenames: {
+    toPNG: (name: string) => `${name}.png`,
+  },
+}));
+
+import { toPng } from "html-to-image";
 import { toast } from "@/components/ui/use-toast";
+import { Logger } from "@/utils/Logger";
 
 describe("withLoadingToast", () => {
   beforeEach(() => {
@@ -30,6 +59,7 @@ describe("withLoadingToast", () => {
     expect(toast).toHaveBeenCalledTimes(1);
     expect(toast).toHaveBeenCalledWith({
       title: "Loading...",
+      duration: Infinity,
     });
     expect(mockDismiss).toHaveBeenCalledTimes(1);
     expect(result).toBe("success");
@@ -94,3 +124,369 @@ describe("withLoadingToast", () => {
     expect(mockDismiss).toHaveBeenCalledTimes(1);
   });
 });
+
+describe("getImageDataUrlForCell", () => {
+  const mockDataUrl = "data:image/png;base64,mockbase64data";
+  let mockElement: HTMLElement;
+
+  beforeEach(() => {
+    vi.clearAllMocks();
+    mockElement = document.createElement("div");
+    mockElement.id = CellOutputId.create("cell-1" as CellId);
+    document.body.append(mockElement);
+  });
+
+  afterEach(() => {
+    mockElement.remove();
+  });
+
+  it("should return undefined if element is not found", async () => {
+    const result = await getImageDataUrlForCell("nonexistent" as CellId);
+
+    expect(result).toBeUndefined();
+    expect(Logger.error).toHaveBeenCalledWith(
+      "Output element not found for cell nonexistent",
+    );
+  });
+
+  it("should capture screenshot and return data URL", async () => {
+    vi.mocked(toPng).mockResolvedValue(mockDataUrl);
+
+    const result = await getImageDataUrlForCell("cell-1" as CellId);
+
+    expect(result).toBe(mockDataUrl);
+    expect(toPng).toHaveBeenCalledWith(mockElement);
+  });
+
+  it("should add printing classes before capture", async () => {
+    vi.mocked(toPng).mockImplementation(async () => {
+      // Check classes are applied during capture
+      expect(mockElement.classList.contains("printing-output")).toBe(true);
+      expect(document.body.classList.contains("printing")).toBe(true);
+      expect(mockElement.style.overflow).toBe("auto");
+      return mockDataUrl;
+    });
+
+    await getImageDataUrlForCell("cell-1" as CellId);
+  });
+
+  it("should remove printing classes after capture", async () => {
+    vi.mocked(toPng).mockResolvedValue(mockDataUrl);
+
+    await getImageDataUrlForCell("cell-1" as CellId);
+
+    expect(mockElement.classList.contains("printing-output")).toBe(false);
+    expect(document.body.classList.contains("printing")).toBe(false);
+  });
+
+  it("should restore original overflow style after capture", async () => {
+    mockElement.style.overflow = "hidden";
+    vi.mocked(toPng).mockResolvedValue(mockDataUrl);
+
+    await getImageDataUrlForCell("cell-1" as CellId);
+
+    expect(mockElement.style.overflow).toBe("hidden");
+  });
+
+  it("should throw error on failure", async () => {
+    vi.mocked(toPng).mockRejectedValue(new Error("Capture failed"));
+
+    await expect(getImageDataUrlForCell("cell-1" as CellId)).rejects.toThrow(
+      "Capture failed",
+    );
+  });
+
+  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(mockElement.classList.contains("printing-output")).toBe(false);
+    expect(document.body.classList.contains("printing")).toBe(false);
+    expect(mockElement.style.overflow).toBe("scroll");
+  });
+
+  it("should maintain body.printing during concurrent captures", async () => {
+    // Create a second element
+    const mockElement2 = document.createElement("div");
+    mockElement2.id = CellOutputId.create("cell-2" as CellId);
+    document.body.append(mockElement2);
+
+    // Track body.printing state during each capture
+    const printingStateDuringCaptures: boolean[] = [];
+    let resolveFirst: () => void;
+    let resolveSecond: () => void;
+
+    const firstPromise = new Promise<void>((resolve) => {
+      resolveFirst = resolve;
+    });
+    const secondPromise = new Promise<void>((resolve) => {
+      resolveSecond = resolve;
+    });
+
+    vi.mocked(toPng).mockImplementation(async (element) => {
+      printingStateDuringCaptures.push(
+        document.body.classList.contains("printing"),
+      );
+
+      // Simulate async work - first capture takes longer
+      await (element.id.includes("cell-1") ? firstPromise : secondPromise);
+
+      // Check state again after waiting
+      printingStateDuringCaptures.push(
+        document.body.classList.contains("printing"),
+      );
+
+      return mockDataUrl;
+    });
+
+    // Start both captures concurrently
+    const capture1 = getImageDataUrlForCell("cell-1" as CellId);
+    const capture2 = getImageDataUrlForCell("cell-2" as CellId);
+
+    // Let second capture complete first
+    resolveSecond!();
+    await new Promise((r) => setTimeout(r, 0));
+
+    // body.printing should still be present because cell-1 is still capturing
+    expect(document.body.classList.contains("printing")).toBe(true);
+
+    // Now let first capture complete
+    resolveFirst!();
+    await Promise.all([capture1, capture2]);
+
+    // After all captures complete, body.printing should be removed
+    expect(document.body.classList.contains("printing")).toBe(false);
+
+    // All captures should have seen body.printing = true
+    expect(printingStateDuringCaptures.every(Boolean)).toBe(true);
+
+    mockElement2.remove();
+  });
+});
+
+describe("downloadHTMLAsImage", () => {
+  const mockDataUrl = "data:image/png;base64,mockbase64data";
+  let mockElement: HTMLElement;
+  let mockAppEl: HTMLElement;
+  let mockAnchor: HTMLAnchorElement;
+
+  beforeEach(() => {
+    vi.clearAllMocks();
+    mockElement = document.createElement("div");
+    mockAppEl = document.createElement("div");
+    mockAppEl.id = "App";
+    // Mock scrollTo since jsdom doesn't implement it
+    mockAppEl.scrollTo = vi.fn();
+    document.body.append(mockElement);
+    document.body.append(mockAppEl);
+
+    // Mock anchor element for download
+    mockAnchor = document.createElement("a");
+    vi.spyOn(document, "createElement").mockReturnValue(mockAnchor);
+    vi.spyOn(mockAnchor, "click").mockImplementation(() => {
+      // <noop></noop>
+    });
+    vi.spyOn(mockAnchor, "remove").mockImplementation(() => {
+      // noop
+    });
+  });
+
+  afterEach(() => {
+    mockElement.remove();
+    mockAppEl.remove();
+    vi.restoreAllMocks();
+  });
+
+  it("should download image without prepare function", async () => {
+    vi.mocked(toPng).mockResolvedValue(mockDataUrl);
+
+    await downloadHTMLAsImage({ element: mockElement, filename: "test" });
+
+    expect(toPng).toHaveBeenCalledWith(mockElement);
+    expect(mockAnchor.href).toBe(mockDataUrl);
+    expect(mockAnchor.download).toBe("test.png");
+    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();
+    const prepare = vi.fn().mockReturnValue(cleanup);
+
+    await downloadHTMLAsImage({
+      element: mockElement,
+      filename: "test",
+      prepare,
+    });
+
+    expect(prepare).toHaveBeenCalledWith(mockElement);
+    expect(cleanup).toHaveBeenCalled();
+  });
+
+  it("should not add body.printing when prepare is provided", async () => {
+    let bodyPrintingDuringCapture = false;
+    vi.mocked(toPng).mockImplementation(async () => {
+      // Capture the state during toPng execution
+      bodyPrintingDuringCapture = document.body.classList.contains("printing");
+      return mockDataUrl;
+    });
+    const cleanup = vi.fn();
+    const prepare = vi.fn().mockReturnValue(cleanup);
+
+    await downloadHTMLAsImage({
+      element: mockElement,
+      filename: "test",
+      prepare,
+    });
+
+    // body.printing should NOT be added by downloadHTMLAsImage when prepare is provided
+    // (the prepare function is responsible for managing its own classes)
+    expect(bodyPrintingDuringCapture).toBe(false);
+    expect(document.body.classList.contains("printing")).toBe(false);
+    expect(prepare).toHaveBeenCalledWith(mockElement);
+    expect(cleanup).toHaveBeenCalled();
+  });
+
+  it("should show error toast on failure", async () => {
+    vi.mocked(toPng).mockRejectedValue(new Error("Failed"));
+
+    await downloadHTMLAsImage({ element: mockElement, filename: "test" });
+
+    expect(toast).toHaveBeenCalledWith({
+      title: "Error",
+      description: "Failed to download as PNG.",
+      variant: "danger",
+    });
+  });
+
+  it("should cleanup on failure", async () => {
+    vi.mocked(toPng).mockRejectedValue(new Error("Failed"));
+
+    await downloadHTMLAsImage({ element: mockElement, filename: "test" });
+
+    expect(document.body.classList.contains("printing")).toBe(false);
+  });
+});
+
+describe("downloadCellOutputAsImage", () => {
+  const mockDataUrl = "data:image/png;base64,mockbase64data";
+  let mockElement: HTMLElement;
+  let mockAppEl: HTMLElement;
+  let mockAnchor: HTMLAnchorElement;
+
+  beforeEach(() => {
+    vi.clearAllMocks();
+    mockElement = document.createElement("div");
+    mockElement.id = CellOutputId.create("cell-1" as CellId);
+    mockAppEl = document.createElement("div");
+    mockAppEl.id = "App";
+    // Mock scrollTo since jsdom doesn't implement it
+    mockAppEl.scrollTo = vi.fn();
+    document.body.append(mockElement);
+    document.body.append(mockAppEl);
+
+    mockAnchor = document.createElement("a");
+    vi.spyOn(document, "createElement").mockReturnValue(mockAnchor);
+    vi.spyOn(mockAnchor, "click").mockImplementation(() => {
+      // <noop></noop>
+    });
+    vi.spyOn(mockAnchor, "remove").mockImplementation(() => {
+      // <noop></noop>
+    });
+  });
+
+  afterEach(() => {
+    mockElement.remove();
+    mockAppEl.remove();
+    vi.restoreAllMocks();
+  });
+
+  it("should return early if element not found", async () => {
+    await downloadCellOutputAsImage("nonexistent" as CellId, "test");
+
+    expect(toPng).not.toHaveBeenCalled();
+    expect(Logger.error).toHaveBeenCalledWith(
+      "Output element not found for cell nonexistent",
+    );
+  });
+
+  it("should download cell output as image", async () => {
+    vi.mocked(toPng).mockResolvedValue(mockDataUrl);
+
+    await downloadCellOutputAsImage("cell-1" as CellId, "result");
+
+    expect(toPng).toHaveBeenCalledWith(mockElement);
+    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.classList.contains("printing-output")).toBe(true);
+      expect(document.body.classList.contains("printing")).toBe(true);
+      expect(mockElement.style.overflow).toBe("auto");
+      return mockDataUrl;
+    });
+
+    await downloadCellOutputAsImage("cell-1" as CellId, "result");
+  });
+
+  it("should cleanup after download", async () => {
+    mockElement.style.overflow = "visible";
+    vi.mocked(toPng).mockResolvedValue(mockDataUrl);
+
+    await downloadCellOutputAsImage("cell-1" as CellId, "result");
+
+    expect(mockElement.classList.contains("printing-output")).toBe(false);
+    expect(document.body.classList.contains("printing")).toBe(false);
+    expect(mockElement.style.overflow).toBe("visible");
+  });
+});
+
+describe("downloadByURL", () => {
+  let mockAnchor: HTMLAnchorElement;
+
+  beforeEach(() => {
+    mockAnchor = document.createElement("a");
+    vi.spyOn(document, "createElement").mockReturnValue(mockAnchor);
+    vi.spyOn(mockAnchor, "click").mockImplementation(() => {
+      // <noop></noop>
+    });
+    vi.spyOn(mockAnchor, "remove").mockImplementation(() => {
+      // <noop></noop>
+    });
+  });
+
+  afterEach(() => {
+    vi.restoreAllMocks();
+  });
+
+  it("should create anchor, set attributes, click, and remove", () => {
+    downloadByURL("data:test", "filename.png");
+
+    expect(document.createElement).toHaveBeenCalledWith("a");
+    expect(mockAnchor.href).toBe("data:test");
+    expect(mockAnchor.download).toBe("filename.png");
+    expect(mockAnchor.click).toHaveBeenCalled();
+    expect(mockAnchor.remove).toHaveBeenCalled();
+  });
+});