@marimo-team/islands 0.23.17-dev1 → 0.23.17-dev4

package/src/components/editor/actions/pair-with-agent-commands.ts

@@ -2,6 +2,7 @@
 
 import { assertNever } from "@/utils/assertNever";
 import { KnownQueryParams } from "@/core/constants";
+import { shellQuote } from "@/utils/shell";
 
 export type AgentTab = "claude" | "codex" | "opencode" | "prompt";
 
@@ -23,27 +24,6 @@ export function getMarimoCommand(): string {
   return import.meta.env.DEV ? "uv run marimo" : "uvx marimo@latest";
 }
 
-/**
- * POSIX-quote a value for safe embedding in a shell command. These commands are
- * meant to be copied into a terminal, so a url/token containing shell
- * metacharacters (`'`, `&`, `$(...)`, ...) must not break out of its argument.
- *
- * Mirrors Python's `shlex.quote` (used on the CLI side in
- * `marimo/_cli/pair/commands.py`) so both sides produce identical commands:
- * values that are already shell-safe are left as-is for readability, and
- * anything else is single-quoted with embedded quotes escaped as `'"'"'`.
- */
-export function shellQuote(value: string): string {
-  if (value === "") {
-    return "''";
-  }
-  // Same "safe" character set as CPython's shlex.quote (ASCII \w plus a few).
-  if (/^[\w@%+=:,./-]+$/.test(value)) {
-    return value;
-  }
-  return `'${value.replaceAll("'", `'"'"'`)}'`;
-}
-
 /** Return the server file key from a page URL, preserving its decoded value. */
 export function getFileFromURL(href: string): string | undefined {
   const file = new URL(href).searchParams.get(KnownQueryParams.filePath);

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

@@ -1,6 +1,7 @@
 /* Copyright 2026 Marimo. All rights reserved. */
 
 import { useAtom, useAtomValue, useSetAtom } from "jotai";
+import type { RefObject } from "react";
 import {
   BookMarkedIcon,
   CheckIcon,
@@ -51,7 +52,6 @@ import { GitHubIcon } from "@/components/icons/github";
 import { MarimoPlusIcon } from "@/components/icons/marimo-icons";
 import { YouTubeIcon } from "@/components/icons/youtube";
 import { useImperativeModal } from "@/components/modal/ImperativeModal";
-import { renderShortcut } from "@/components/shortcuts/renderShortcut";
 import { PairWithAgentModal } from "@/components/editor/actions/pair-with-agent-modal";
 import { ShareStaticNotebookModal } from "@/components/static-html/share-modal";
 import { toast } from "@/components/ui/use-toast";
@@ -66,30 +66,15 @@ import { disabledCellIds } from "@/core/cells/utils";
 import { capabilitiesAtom } from "@/core/config/capabilities";
 import { aiEnabledAtom, useResolvedMarimoConfig } from "@/core/config/config";
 import { Constants } from "@/core/constants";
-import {
-  updateCellOutputsWithScreenshots,
-  useEnrichCellOutputs,
-} from "@/core/export/hooks";
 import { useLayoutActions, useLayoutState } from "@/core/layout/layout";
 import { useTogglePresenting } from "@/core/layout/useTogglePresenting";
 import { kioskModeAtom, viewStateAtom } from "@/core/mode";
 import { useRequestClient } from "@/core/network/requests";
 import { useFilename } from "@/core/saving/filename";
-import { downloadAsHTML } from "@/core/static/download-html";
 import { createShareableLink } from "@/core/wasm/share";
 import { isWasm } from "@/core/wasm/utils";
 import { copyToClipboard } from "@/utils/copy";
-import {
-  ADD_PRINTING_CLASS,
-  downloadAsPDF,
-  downloadBlob,
-  downloadExportedFile,
-  downloadHTMLAsImage,
-  withLoadingToast,
-} from "@/utils/download";
-import { Filenames } from "@/utils/filenames";
 import { Objects } from "@/utils/objects";
-import type { ProgressState } from "@/utils/progress";
 import { Strings } from "@/utils/strings";
 import { newNotebookURL } from "@/utils/urls";
 import { useRunAllCells } from "../cell/useRunCells";
@@ -100,7 +85,13 @@ import { keyboardShortcutsAtom } from "../controls/keyboard-shortcuts";
 import { commandPaletteAtom } from "../controls/state";
 import { displayLayoutName, getLayoutIcon } from "../renderers/layout-select";
 import { LAYOUT_TYPES } from "../renderers/types";
-import { runServerSidePDFDownload } from "./pdf-export";
+import { ExportDialog } from "./export-dialog/export-dialog";
+import {
+  applyExportOptionOverrides,
+  exportOptionsAtom,
+  type ExportFormat,
+  type ExportOptionOverrides,
+} from "./export-dialog/state";
 import type { ActionButton } from "./types";
 import { useCopyNotebook } from "./useCopyNotebook";
 import { useRestartKernel } from "./useRestartKernel";
@@ -111,7 +102,11 @@ const NOOP_HANDLER = (event?: Event) => {
   event?.stopPropagation();
 };
 
-export function useNotebookActions() {
+export function useNotebookActions({
+  exportDialogReturnFocusRef,
+}: {
+  exportDialogReturnFocusRef?: RefObject<HTMLElement | null>;
+} = {}) {
   const filename = useFilename();
   const { openModal, closeModal } = useImperativeModal();
   const { toggleApplication } = useChromeActions();
@@ -138,15 +133,8 @@ export function useNotebookActions() {
   const setSettingsDialogOpen = useSetAtom(settingDialogAtom);
   const { handleClick: openSettings } = useOpenSettingsToTab();
   const setKeyboardShortcutsOpen = useSetAtom(keyboardShortcutsAtom);
-  const {
-    exportAsIPYNB,
-    exportAsMarkdown,
-    exportAsScript,
-    readCode,
-    saveCellConfig,
-    updateCellOutputs,
-  } = useRequestClient();
-  const takeScreenshots = useEnrichCellOutputs();
+  const setExportOptions = useSetAtom(exportOptionsAtom);
+  const { readCode, saveCellConfig } = useRequestClient();
 
   const hasDisabledCells = useAtomValue(hasDisabledCellsAtom);
   const canUndoDeletes = useAtomValue(canUndoDeletesAtom);
@@ -158,10 +146,6 @@ export function useNotebookActions() {
   const sharingHtmlEnabled = resolvedConfig.sharing?.html ?? true;
   const sharingWasmEnabled = resolvedConfig.sharing?.wasm ?? true;
   const sharingMolabEnabled = resolvedConfig.sharing?.molab ?? true;
-
-  // Server-side PDF export is always available outside WASM.
-  // Browser print fallback is used in WASM.
-  const serverSidePdfEnabled = !isWasm();
   const isSlidesLayout = selectedLayout === "slides";
 
   const renderCheckboxElement = (checked: boolean) => (
@@ -170,82 +154,39 @@ export function useNotebookActions() {
     </div>
   );
 
-  const renderRecommendedElement = (recommended: boolean) => {
-    if (!recommended) {
-      return null;
+  const openExportDialog = (
+    initialFormat?: ExportFormat,
+    optionOverrides?: ExportOptionOverrides,
+  ) => {
+    if (optionOverrides) {
+      setExportOptions((current) =>
+        applyExportOptionOverrides(current, optionOverrides),
+      );
     }
-    return (
-      <span className="ml-3 shrink-0 rounded-full border border-emerald-200 bg-emerald-50 px-2 py-0.5 text-[10px] font-semibold uppercase tracking-wide text-emerald-700">
-        Recommended
-      </span>
+    openModal(
+      <ExportDialog
+        initialFormat={initialFormat}
+        onClose={closeModal}
+        returnFocusRef={exportDialogReturnFocusRef}
+      />,
     );
   };
 
-  const downloadServerSidePDF = async ({
-    preset,
-    title,
-  }: {
-    preset: "document" | "slides";
-    title: string;
-  }) => {
-    if (!filename) {
-      toastNotebookMustBeNamed();
-      return;
-    }
-
-    const runDownload = async (progress: ProgressState) => {
-      await runServerSidePDFDownload({
-        exportOptions: {
-          webpdf: false,
-          preset,
-          includeInputs: true,
-          includeOutputs: true,
-        },
-        captureOutputs: () =>
-          updateCellOutputsWithScreenshots({
-            takeScreenshots: () => takeScreenshots({ progress }),
-            updateCellOutputs,
-          }),
-        downloadPDF: downloadAsPDF,
-      });
-    };
-    await withLoadingToast(title, runDownload);
-  };
-
-  const handleDocumentPDF = async () => {
-    if (serverSidePdfEnabled) {
-      await downloadServerSidePDF({
-        preset: "document",
-        title: "Downloading Document PDF...",
-      });
-      return;
-    }
-    const beforeprint = new Event("export-beforeprint");
-    const afterprint = new Event("export-afterprint");
-    window.dispatchEvent(beforeprint);
-    setTimeout(() => window.print(), 0);
-    setTimeout(() => window.dispatchEvent(afterprint), 0);
-  };
-
-  const handleDownloadAsIPYNB = async () => {
-    if (!filename) {
-      toastNotebookMustBeNamed();
-      return;
-    }
-
-    const runDownload = async (progress: ProgressState) => {
-      await updateCellOutputsWithScreenshots({
-        takeScreenshots: () => takeScreenshots({ progress }),
-        updateCellOutputs,
-      });
-      const exportedFile = await exportAsIPYNB({ download: false });
-      downloadExportedFile(exportedFile);
-    };
-
-    await withLoadingToast("Downloading IPYNB...", runDownload);
-  };
-
   const actions: ActionButton[] = [
+    {
+      icon: <DownloadIcon size={14} strokeWidth={1.5} />,
+      label: "Export…",
+      additionalKeywords: [
+        "download",
+        "html",
+        "markdown",
+        "ipynb",
+        "pdf",
+        "script",
+        "png",
+      ],
+      handle: () => openExportDialog(),
+    },
     {
       icon: <DownloadIcon size={14} strokeWidth={1.5} />,
       label: "Download",
@@ -254,83 +195,44 @@ export function useNotebookActions() {
         {
           icon: <FolderDownIcon size={14} strokeWidth={1.5} />,
           label: "Download as HTML",
-          handle: async () => {
-            if (!filename) {
-              toastNotebookMustBeNamed();
-              return;
-            }
-            await downloadAsHTML({ includeCode: true });
-          },
+          handle: () =>
+            openExportDialog("html", { html: { includeCode: true } }),
         },
         {
           icon: <FolderDownIcon size={14} strokeWidth={1.5} />,
           label: "Download as HTML (exclude code)",
-          handle: async () => {
-            if (!filename) {
-              toastNotebookMustBeNamed();
-              return;
-            }
-            await downloadAsHTML({ includeCode: false });
-          },
+          handle: () =>
+            openExportDialog("html", { html: { includeCode: false } }),
         },
         {
           icon: (
             <MarkdownIcon strokeWidth={1.5} style={{ width: 14, height: 14 }} />
           ),
           label: "Download as Markdown",
-          handle: async () => {
-            const exportedFile = await exportAsMarkdown({ download: false });
-            downloadExportedFile(exportedFile);
-          },
+          handle: () => openExportDialog("markdown"),
         },
         {
           icon: <NotebookIcon size={14} strokeWidth={1.5} />,
           label: "Download as ipynb",
-          handle: handleDownloadAsIPYNB,
+          handle: () => openExportDialog("ipynb"),
         },
         {
           icon: <CodeIcon size={14} strokeWidth={1.5} />,
           label: "Download notebook source",
-          handle: async () => {
-            const code = await readCode();
-            downloadBlob(
-              new Blob([code.contents], { type: "text/plain" }),
-              Filenames.toPY(document.title),
-            );
-          },
+          handle: () =>
+            openExportDialog("script", { script: { type: "source" } }),
         },
         {
           icon: <CodeIcon size={14} strokeWidth={1.5} />,
           label: "Download flat script",
-          handle: async () => {
-            const exportedFile = await exportAsScript({ download: false });
-            downloadExportedFile(exportedFile);
-          },
+          handle: () =>
+            openExportDialog("script", { script: { type: "flat" } }),
         },
         {
           divider: true,
           icon: <ImageIcon size={14} strokeWidth={1.5} />,
           label: "Download as PNG",
-          disabled: viewState.mode !== "present",
-          tooltip:
-            viewState.mode === "present" ? undefined : (
-              <span>
-                Only available in app view. <br />
-                Toggle with: {renderShortcut("global.hideCode", false)}
-              </span>
-            ),
-          handle: async () => {
-            const app = document.getElementById("App");
-            if (!app) {
-              return;
-            }
-            await downloadHTMLAsImage({
-              element: app,
-              filename: document.title,
-              // Add body.printing ONLY when converting the whole notebook to a screenshot
-              prepare: ADD_PRINTING_CLASS,
-            });
-          },
+          handle: () => openExportDialog("png"),
         },
         isSlidesLayout
           ? {
@@ -342,19 +244,24 @@ export function useNotebookActions() {
                 {
                   icon: <FileIcon size={14} strokeWidth={1.5} />,
                   label: "Document Layout",
-                  handle: handleDocumentPDF,
+                  handle: () =>
+                    openExportDialog("pdf", {
+                      pdf: { preset: "document" },
+                    }),
                 },
                 {
                   icon: <FileIcon size={14} strokeWidth={1.5} />,
                   label: "Slides Layout",
-                  rightElement: renderRecommendedElement(true),
-                  hidden: !serverSidePdfEnabled,
-                  handle: async () => {
-                    await downloadServerSidePDF({
-                      preset: "slides",
-                      title: "Downloading Slides PDF...",
-                    });
-                  },
+                  hidden: isWasm(),
+                  rightElement: (
+                    <span className="ml-3 shrink-0 rounded-full border border-emerald-200 bg-emerald-50 px-2 py-0.5 text-[10px] font-semibold uppercase tracking-wide text-emerald-700">
+                      Recommended
+                    </span>
+                  ),
+                  handle: () =>
+                    openExportDialog("pdf", {
+                      pdf: { preset: "slides" },
+                    }),
                 },
               ],
             }
@@ -362,7 +269,10 @@ export function useNotebookActions() {
               divider: true,
               icon: <FileIcon size={14} strokeWidth={1.5} />,
               label: "Download as PDF",
-              handle: handleDocumentPDF,
+              handle: () =>
+                openExportDialog("pdf", {
+                  pdf: { preset: "document" },
+                }),
             },
       ],
     },
@@ -740,11 +650,3 @@ export function useNotebookActions() {
       return action;
     });
 }
-
-function toastNotebookMustBeNamed() {
-  toast({
-    title: "Error",
-    description: "Notebooks must be named to be exported.",
-    variant: "danger",
-  });
-}

package/src/components/editor/connections/add-connection-dialog.tsx

@@ -87,7 +87,7 @@ export const AddConnectionDialogContent: React.FC<{
           <span className="block">{codeSnippetHint}</span>
         </DialogDescription>
       </DialogHeader>
-      <AutoDiscoveredDataSources onSubmit={onClose} />
+      <AutoDiscoveredDataSources onSubmit={onClose} className="-mt-2" />
       <Tabs
         value={activeTab}
         onValueChange={(v) => setActiveTab(v as ConnectionTab)}

package/src/components/editor/connections/quick-add-data-sources.tsx

@@ -5,11 +5,13 @@ import { Tooltip, TooltipProvider } from "@/components/ui/tooltip";
 import type { DetectedDataSource } from "@/core/datasets/data-source-discovery";
 import { useDataSourceDiscovery } from "@/hooks/useDataSourceDiscovery";
 import { useInsertCode } from "./components";
+import { cn } from "@/utils/cn";
 
 export const QuickAddDataSources: React.FC<{
+  className?: string;
   sources: DetectedDataSource[];
   onAdd: (source: DetectedDataSource) => void;
-}> = ({ sources, onAdd }) => {
+}> = ({ className, sources, onAdd }) => {
   if (sources.length === 0) {
     return null;
   }
@@ -17,12 +19,15 @@ export const QuickAddDataSources: React.FC<{
   return (
     <section
       aria-labelledby="quick-add-data-sources-title"
-      className="rounded-md border bg-muted/30 px-3 py-2"
+      className={cn(
+        "rounded-full bg-[linear-gradient(135deg,var(--blue-2),var(--purple-3))] px-3 py-2",
+        className,
+      )}
     >
       <div className="flex flex-wrap items-center gap-2">
         <div className="mr-1 flex items-center gap-1.5">
-          <SparklesIcon className="h-3.5 w-3.5 text-muted-foreground" />
-          <h3 id="quick-add-data-sources-title" className="text-sm font-medium">
+          <SparklesIcon className="h-3.5 w-3.5 text-(--blue-9)" />
+          <h3 id="quick-add-data-sources-title" className="text-sm">
             Quick add
           </h3>
         </div>
@@ -36,10 +41,10 @@ export const QuickAddDataSources: React.FC<{
               <button
                 type="button"
                 aria-label={`Add ${source.displayName} connection`}
-                className="inline-flex items-center gap-1 rounded-full border border-(--blue-8) bg-(--blue-2) px-2.5 py-1 text-xs font-semibold text-(--blue-11) transition-colors hover:bg-(--blue-3) focus:outline-hidden focus:ring-2 focus:ring-ring focus:ring-offset-2"
+                className="inline-flex items-center gap-1 rounded-full border border-border/60 bg-background px-2.5 py-1 text-xs font-medium text-foreground/90 transition-colors hover:border-(--blue-7) focus:outline-hidden focus:ring-2 focus:ring-ring focus:ring-offset-2"
                 onClick={() => onAdd(source)}
               >
-                <PlusIcon className="h-3 w-3" />
+                <PlusIcon className="h-3 w-3 text-muted-foreground" />
                 {source.displayName}
               </button>
             </Tooltip>
@@ -90,12 +95,14 @@ const DetectedDataSourceDetails: React.FC<{
 
 export const AutoDiscoveredDataSources: React.FC<{
   onSubmit: () => void;
-}> = ({ onSubmit }) => {
+  className?: string;
+}> = ({ onSubmit, className }) => {
   const insertCode = useInsertCode();
   const { data } = useDataSourceDiscovery();
 
   return (
     <QuickAddDataSources
+      className={className}
       sources={data ?? []}
       onAdd={(source) => {
         insertCode(source.code);

package/src/components/editor/controls/__tests__/notebook-menu-dropdown.test.tsx

@@ -0,0 +1,187 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+import { fireEvent, render, screen, waitFor } from "@testing-library/react";
+import { Provider } from "jotai";
+import type React from "react";
+import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
+import { MockRequestClient } from "@/__mocks__/requests";
+import { ModalProvider } from "@/components/modal/ImperativeModal";
+import { TooltipProvider } from "@/components/ui/tooltip";
+import { layoutStateAtom } from "@/core/layout/layout";
+import { kioskModeAtom, viewStateAtom } from "@/core/mode";
+import { requestClientAtom } from "@/core/network/requests";
+import { filenameAtom } from "@/core/saving/file-state";
+import { store } from "@/core/state/jotai";
+import { isWasm } from "@/core/wasm/utils";
+import {
+  DEFAULT_EXPORT_OPTIONS,
+  exportOptionsAtom,
+  lastExportFormatAtom,
+} from "../../actions/export-dialog/state";
+import { NotebookMenuDropdown } from "../notebook-menu-dropdown";
+
+vi.mock("@/core/wasm/utils", async (importOriginal) => {
+  const actual = await importOriginal<typeof import("@/core/wasm/utils")>();
+  return { ...actual, isWasm: vi.fn(() => false) };
+});
+
+function wrapper({ children }: { children: React.ReactNode }) {
+  return (
+    <Provider store={store}>
+      <TooltipProvider>
+        <ModalProvider>{children}</ModalProvider>
+      </TooltipProvider>
+    </Provider>
+  );
+}
+
+function openNotebookMenu() {
+  fireEvent.pointerDown(screen.getByTestId("notebook-menu-dropdown"), {
+    button: 0,
+    ctrlKey: false,
+  });
+}
+
+async function openDownloadMenu() {
+  openNotebookMenu();
+  fireEvent.click(await screen.findByRole("menuitem", { name: "Download" }));
+}
+
+async function selectDownload(name: string | RegExp) {
+  await openDownloadMenu();
+  fireEvent.click(await screen.findByRole("menuitem", { name }));
+}
+
+describe("NotebookMenuDropdown", () => {
+  beforeEach(() => {
+    vi.clearAllMocks();
+    localStorage.clear();
+    store.set(requestClientAtom, MockRequestClient.create());
+    store.set(filenameAtom, "/project/notebook.py");
+    store.set(viewStateAtom, { mode: "edit", cellAnchor: null });
+    store.set(kioskModeAtom, false);
+    store.set(layoutStateAtom, {
+      selectedLayout: "vertical",
+      layoutData: {},
+    });
+    store.set(exportOptionsAtom, DEFAULT_EXPORT_OPTIONS);
+    store.set(lastExportFormatAtom, "html");
+    vi.mocked(isWasm).mockReturnValue(false);
+    vi.stubGlobal("PointerEvent", MouseEvent);
+    vi.stubGlobal("matchMedia", () => ({
+      matches: false,
+      addEventListener: vi.fn(),
+      removeEventListener: vi.fn(),
+    }));
+  });
+
+  afterEach(() => {
+    document.title = "";
+    vi.unstubAllGlobals();
+  });
+
+  it("returns focus to the notebook menu when the dialog closes", async () => {
+    render(<NotebookMenuDropdown />, { wrapper });
+    const menuButton = screen.getByTestId("notebook-menu-dropdown");
+
+    openNotebookMenu();
+    fireEvent.click(
+      await screen.findByRole("menuitem", {
+        name: "Export…",
+      }),
+    );
+    fireEvent.click(await screen.findByRole("button", { name: "Close" }));
+
+    await waitFor(() => expect(menuButton).toHaveFocus());
+  });
+
+  it("opens the HTML shortcut with code excluded", async () => {
+    render(<NotebookMenuDropdown />, { wrapper });
+
+    await selectDownload("Download as HTML (exclude code)");
+
+    expect(await screen.findByTestId("export-dialog")).toBeVisible();
+    expect(screen.getByTestId("export-format-html")).toHaveAttribute(
+      "aria-selected",
+      "true",
+    );
+    expect(
+      screen.getByRole("switch", { name: "Include code" }),
+    ).not.toBeChecked();
+  });
+
+  it("opens the slides PDF shortcut with the slides layout", async () => {
+    store.set(layoutStateAtom, {
+      selectedLayout: "slides",
+      layoutData: {},
+    });
+    render(<NotebookMenuDropdown />, { wrapper });
+
+    await openDownloadMenu();
+    fireEvent.click(
+      await screen.findByRole("menuitem", {
+        name: "Download as PDF",
+      }),
+    );
+    fireEvent.click(
+      await screen.findByRole("menuitem", {
+        name: /Slides Layout/,
+      }),
+    );
+
+    expect(await screen.findByTestId("export-dialog")).toBeVisible();
+    expect(screen.getByTestId("export-format-pdf")).toHaveAttribute(
+      "aria-selected",
+      "true",
+    );
+    expect(screen.getByRole("radio", { name: "Slides" })).toBeChecked();
+  });
+
+  it("hides the slides PDF shortcut in WebAssembly", async () => {
+    vi.mocked(isWasm).mockReturnValue(true);
+    store.set(layoutStateAtom, {
+      selectedLayout: "slides",
+      layoutData: {},
+    });
+    render(<NotebookMenuDropdown />, { wrapper });
+
+    await openDownloadMenu();
+    fireEvent.click(
+      await screen.findByRole("menuitem", {
+        name: "Download as PDF",
+      }),
+    );
+
+    expect(
+      screen.getByRole("menuitem", { name: "Document Layout" }),
+    ).toBeVisible();
+    expect(
+      screen.queryByRole("menuitem", { name: /Slides Layout/ }),
+    ).not.toBeInTheDocument();
+  });
+
+  it("preselects each Python format from its download shortcut", async () => {
+    render(<NotebookMenuDropdown />, { wrapper });
+
+    await selectDownload("Download notebook source");
+
+    expect(await screen.findByTestId("export-dialog")).toBeVisible();
+    expect(screen.getByTestId("export-format-script")).toHaveAttribute(
+      "aria-selected",
+      "true",
+    );
+    expect(
+      screen.getByRole("radio", { name: "Notebook source" }),
+    ).toBeChecked();
+
+    fireEvent.click(screen.getByRole("button", { name: "Close" }));
+    await waitFor(() =>
+      expect(screen.queryByTestId("export-dialog")).not.toBeInTheDocument(),
+    );
+
+    await selectDownload("Download flat script");
+
+    expect(await screen.findByTestId("export-dialog")).toBeVisible();
+    expect(screen.getByRole("radio", { name: "Flat script" })).toBeChecked();
+  });
+});

package/src/components/editor/controls/notebook-menu-dropdown.tsx

@@ -1,7 +1,7 @@
 /* Copyright 2026 Marimo. All rights reserved. */
 
 import { MenuIcon } from "lucide-react";
-import React from "react";
+import React, { useRef } from "react";
 import { useLocale } from "react-aria";
 import { Button } from "@/components/editor/inputs/Inputs";
 import {
@@ -33,7 +33,8 @@ export const NotebookMenuDropdown: React.FC<Props> = ({
   disabled = false,
   tooltip = "Actions",
 }) => {
-  const actions = useNotebookActions();
+  const exportDialogReturnFocusRef = useRef<HTMLButtonElement>(null);
+  const actions = useNotebookActions({ exportDialogReturnFocusRef });
   const { locale } = useLocale();
   // Create tooltip content with keyboard shortcut decoration
   const tooltipContent = (
@@ -50,6 +51,7 @@ export const NotebookMenuDropdown: React.FC<Props> = ({
 
   const button = (
     <Button
+      ref={exportDialogReturnFocusRef}
       aria-label="Config"
       shape="circle"
       size="small"