@marimo-team/islands 0.19.5-dev4 → 0.19.5-dev43

package/src/components/app-config/user-config-form.tsx

@@ -50,6 +50,7 @@ import { Banner } from "@/plugins/impl/common/error-banner";
 import { THEMES } from "@/theme/useTheme";
 import { arrayToggle } from "@/utils/arrays";
 import { cn } from "@/utils/cn";
+import { autoPopulateModels } from "../ai/ai-utils";
 import { keyboardShortcutsAtom } from "../editor/controls/keyboard-shortcuts";
 import { Badge } from "../ui/badge";
 import { ExternalLink } from "../ui/links";
@@ -180,6 +181,28 @@ export const UserConfigForm: React.FC = () => {
     defaultValues: config,
   });
 
+  const setAiModels = (values: UserConfig, dirtyAiConfig: UserConfig["ai"]) => {
+    const { chatModel, editModel } = autoPopulateModels(values);
+    if (chatModel || editModel) {
+      dirtyAiConfig = {
+        ...dirtyAiConfig,
+        models: {
+          ...dirtyAiConfig?.models,
+          ...(chatModel && { chat_model: chatModel }),
+          ...(editModel && { edit_model: editModel }),
+        },
+      } as typeof dirtyAiConfig;
+      if (chatModel) {
+        form.setValue("ai.models.chat_model", chatModel);
+      }
+      if (editModel) {
+        form.setValue("ai.models.edit_model", editModel);
+      }
+    }
+
+    return dirtyAiConfig;
+  };
+
   const onSubmitNotDebounced = async (values: UserConfig) => {
     // Only send values that were actually changed to avoid
     // overwriting backend values the form doesn't manage
@@ -187,6 +210,12 @@ export const UserConfigForm: React.FC = () => {
     if (Object.keys(dirtyValues).length === 0) {
       return; // Nothing changed
     }
+
+    // Auto-populate AI models when credentials are set, makes it easier to get started
+    if (dirtyValues.ai) {
+      dirtyValues.ai = setAiModels(values, dirtyValues.ai);
+    }
+
     await saveUserConfig({ config: dirtyValues }).then(() => {
       // Update local state with form values
       setConfig((prev) => ({ ...prev, ...values }));
@@ -993,7 +1022,7 @@ export const UserConfigForm: React.FC = () => {
                       <br />
                       <br />
                       Running marimo in a{" "}
-                      <ExternalLink href="https://docs.marimo.io/guides/editor_features/package_management.html#running-marimo-in-a-sandbox-environment-uv-only">
+                      <ExternalLink href="https://docs.marimo.io/guides/package_management/inlining_dependencies.html">
                         sandboxed environment
                       </ExternalLink>{" "}
                       is only supported by <Kbd className="inline">uv</Kbd>
@@ -1283,6 +1312,39 @@ export const UserConfigForm: React.FC = () => {
                 </div>
               )}
             />
+            <FormField
+              control={form.control}
+              name="experimental.server_side_pdf_export"
+              render={({ field }) => (
+                <div className="flex flex-col gap-y-1">
+                  <FormItem className={formItemClasses}>
+                    <FormLabel className="font-normal">
+                      Better PDF Export
+                    </FormLabel>
+                    <FormControl>
+                      <Checkbox
+                        data-testid="server-side-pdf-export-checkbox"
+                        checked={field.value === true}
+                        onCheckedChange={field.onChange}
+                      />
+                    </FormControl>
+                  </FormItem>
+                  <IsOverridden
+                    userConfig={config}
+                    name="experimental.server_side_pdf_export"
+                  />
+                  <FormDescription>
+                    Enable PDF export using{" "}
+                    <Kbd className="inline">nbconvert</Kbd> and{" "}
+                    <Kbd className="inline">playwright</Kbd>. Refer to{" "}
+                    <ExternalLink href="https://docs.marimo.io/guides/exporting/#exporting-to-pdf-slides-or-rst">
+                      the docs
+                    </ExternalLink>
+                    .
+                  </FormDescription>
+                </div>
+              )}
+            />
           </SettingGroup>
         );
     }

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

@@ -262,7 +262,7 @@ const ChatInputFooter: React.FC<ChatInputFooterProps> = memo(
 
     return (
       <TooltipProvider>
-        <div className="px-3 py-2 border-t border-border/20 flex flex-row items-center justify-between">
+        <div className="px-3 py-2 border-t border-border/20 flex flex-row flex-wrap items-center justify-between gap-1">
           <div className="flex items-center gap-2">
             <FeatureFlagged feature="chat_modes">
               <Select value={currentMode} onValueChange={saveModeChange}>
@@ -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/data-table/cell-utils.ts

@@ -1,5 +1,7 @@
 /* Copyright 2026 Marimo. All rights reserved. */
 
+import type { CellId } from "@/core/cells/ids";
+
 export const DATA_CELL_ID = "data-cell-id";
 
 export function getCellDomProps(cellId: string) {
@@ -7,3 +9,11 @@ export function getCellDomProps(cellId: string) {
     [DATA_CELL_ID]: cellId,
   };
 }
+
+export const DATA_FOR_CELL_ID = "data-for-cell-id";
+
+export function getCellForDomProps(cellId: CellId) {
+  return {
+    [DATA_FOR_CELL_ID]: cellId,
+  };
+}

package/src/components/editor/Output.tsx

@@ -20,6 +20,7 @@ import { TextOutput } from "./output/TextOutput";
 import { VideoOutput } from "./output/VideoOutput";
 
 import "./output/Outputs.css";
+import { useAtomValue } from "jotai";
 import {
   ChevronsDownUpIcon,
   ChevronsUpDownIcon,
@@ -27,6 +28,7 @@ import {
 } from "lucide-react";
 import { tooltipHandler } from "@/components/charts/tooltip";
 import { useExpandedOutput } from "@/core/cells/outputs";
+import { viewStateAtom } from "@/core/mode";
 import { useIframeCapabilities } from "@/hooks/useIframeCapabilities";
 import { renderHTML } from "@/plugins/core/RenderHTML";
 import { Banner } from "@/plugins/impl/common/error-banner";
@@ -46,10 +48,9 @@ import { renderMimeIcon } from "./renderMimeIcon";
 
 const METADATA_KEY = "__metadata__";
 
-type MimeBundleWithoutMetadata = Record<
-  OutputMessage["mimetype"],
-  { [key: string]: unknown }
->;
+export type MimeType = OutputMessage["mimetype"];
+
+type MimeBundleWithoutMetadata = Record<MimeType, { [key: string]: unknown }>;
 
 type MimeBundle = MimeBundleWithoutMetadata & {
   [METADATA_KEY]?: Record<string, { width?: number; height?: number }>;
@@ -70,7 +71,7 @@ export const OutputRenderer: React.FC<{
   onRefactorWithAI?: OnRefactorWithAI;
   wrapText?: boolean;
   metadata?: { width?: number; height?: number };
-  renderFallback?: (mimetype: OutputMessage["mimetype"]) => React.ReactNode;
+  renderFallback?: (mimetype: MimeType) => React.ReactNode;
 }> = memo((props) => {
   const {
     message,
@@ -90,6 +91,8 @@ export const OutputRenderer: React.FC<{
       case "application/vnd.marimo+mimebundle":
       case "application/vnd.vegalite.v5+json":
       case "application/vnd.vega.v5+json":
+      case "application/vnd.vegalite.v6+json":
+      case "application/vnd.vega.v6+json":
         return typeof data === "string" ? JSON.parse(data) : data;
       default:
         return;
@@ -199,6 +202,8 @@ export const OutputRenderer: React.FC<{
       );
     case "application/vnd.vegalite.v5+json":
     case "application/vnd.vega.v5+json":
+    case "application/vnd.vegalite.v6+json":
+    case "application/vnd.vega.v6+json":
       return (
         <Suspense fallback={<ChartLoadingState />}>
           <LazyVegaEmbed
@@ -216,9 +221,7 @@ export const OutputRenderer: React.FC<{
       return (
         <MimeBundleOutputRenderer
           channel={channel}
-          data={
-            parsedJsonData as Record<OutputMessage["mimetype"], OutputMessage>
-          }
+          data={parsedJsonData as Record<MimeType, OutputMessage>}
         />
       );
     case "application/vnd.jupyter.widget-view+json":
@@ -258,6 +261,8 @@ const MimeBundleOutputRenderer: React.FC<{
   cellId?: CellId;
 }> = memo(({ data, channel, cellId }) => {
   const mimebundle = Array.isArray(data) ? data[0] : data;
+  const { mode } = useAtomValue(viewStateAtom);
+  const appView = mode === "present" || mode === "read";
 
   // Extract metadata if present (e.g., for retina image rendering)
   const metadata = mimebundle[METADATA_KEY];
@@ -265,10 +270,7 @@ const MimeBundleOutputRenderer: React.FC<{
   // Filter out metadata from the mime entries and type narrow
   const mimeEntries = Objects.entries(mimebundle as Record<string, unknown>)
     .filter(([key]) => key !== METADATA_KEY)
-    .map(
-      ([mime, data]) =>
-        [mime, data] as [OutputMessage["mimetype"], CellOutput["data"]],
-    );
+    .map(([mime, data]) => [mime, data] as [MimeType, CellOutput["data"]]);
 
   // If there is none, return null
   const first = mimeEntries[0]?.[0];
@@ -302,7 +304,12 @@ const MimeBundleOutputRenderer: React.FC<{
   return (
     <Tabs defaultValue={first} orientation="vertical">
       <div className="flex">
-        <TabsList className="self-start max-h-none flex flex-col gap-2 mr-4 shrink-0">
+        <TabsList
+          className={cn(
+            "self-start max-h-none flex flex-col gap-2 mr-3 shrink-0",
+            appView && "mt-4",
+          )}
+        >
           {mimeEntries.map(([mime]) => (
             <TabsTrigger
               key={mime}
@@ -315,7 +322,7 @@ const MimeBundleOutputRenderer: React.FC<{
             </TabsTrigger>
           ))}
         </TabsList>
-        <div className="flex-1">
+        <div className="flex-1 w-full">
           {mimeEntries.map(([mime, output]) => (
             <TabsContent key={mime} value={mime}>
               <ErrorBoundary>

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

@@ -54,7 +54,12 @@ import {
 } from "@/core/cells/cells";
 import { disabledCellIds } from "@/core/cells/utils";
 import { useResolvedMarimoConfig } from "@/core/config/config";
+import { getFeatureFlag } from "@/core/config/feature-flag";
 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";
@@ -64,7 +69,12 @@ 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 { downloadBlob, downloadHTMLAsImage } from "@/utils/download";
+import {
+  downloadAsPDF,
+  downloadBlob,
+  downloadHTMLAsImage,
+  withLoadingToast,
+} from "@/utils/download";
 import { Filenames } from "@/utils/filenames";
 import { Objects } from "@/utils/objects";
 import { newNotebookURL } from "@/utils/urls";
@@ -110,18 +120,26 @@ export function useNotebookActions() {
   const setCommandPaletteOpen = useSetAtom(commandPaletteAtom);
   const setSettingsDialogOpen = useSetAtom(settingDialogAtom);
   const setKeyboardShortcutsOpen = useSetAtom(keyboardShortcutsAtom);
-  const { exportAsMarkdown, readCode, saveCellConfig } = useRequestClient();
+  const { exportAsMarkdown, readCode, saveCellConfig, updateCellOutputs } =
+    useRequestClient();
 
   const hasDisabledCells = useAtomValue(hasDisabledCellsAtom);
   const canUndoDeletes = useAtomValue(canUndoDeletesAtom);
   const { selectedLayout } = useLayoutState();
   const { setLayoutView } = useLayoutActions();
   const togglePresenting = useTogglePresenting();
+  const takeScreenshots = useEnrichCellOutputs();
 
   // Fallback: if sharing is undefined, both are enabled by default
   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} />}
@@ -139,11 +157,7 @@ export function useNotebookActions() {
           label: "Download as HTML",
           handle: async () => {
             if (!filename) {
-              toast({
-                variant: "danger",
-                title: "Error",
-                description: "Notebooks must be named to be exported.",
-              });
+              toastNotebookMustBeNamed();
               return;
             }
             await downloadAsHTML({ filename, includeCode: true });
@@ -154,11 +168,7 @@ export function useNotebookActions() {
           label: "Download as HTML (exclude code)",
           handle: async () => {
             if (!filename) {
-              toast({
-                variant: "danger",
-                title: "Error",
-                description: "Notebooks must be named to be exported.",
-              });
+              toastNotebookMustBeNamed();
               return;
             }
             await downloadAsHTML({ filename, includeCode: false });
@@ -205,21 +215,43 @@ 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 (isServerSidePdfExportEnabled) {
+              if (!filename) {
+                toastNotebookMustBeNamed();
+                return;
+              }
+
+              const downloadPDF = async () => {
+                await updateCellOutputsWithScreenshots(
+                  takeScreenshots,
+                  updateCellOutputs,
+                );
+                await downloadAsPDF({
+                  filename: filename,
+                  webpdf: false,
+                });
+              };
+              await withLoadingToast("Downloading PDF...", downloadPDF);
+              return;
+            }
+
             const beforeprint = new Event("export-beforeprint");
             const afterprint = new Event("export-afterprint");
             function print() {
@@ -527,3 +559,11 @@ export function useNotebookActions() {
       return action;
     });
 }
+
+function toastNotebookMustBeNamed() {
+  toast({
+    title: "Error",
+    description: "Notebooks must be named to be exported.",
+    variant: "danger",
+  });
+}

package/src/components/editor/cell/cell-actions.tsx

@@ -12,6 +12,7 @@ import React, {
   useState,
 } from "react";
 import useEvent from "react-use-event-hook";
+import { getCellForDomProps } from "@/components/data-table/cell-utils";
 import {
   renderMinimalShortcut,
   renderShortcut,
@@ -104,7 +105,11 @@ const CellActionsDropdownInternal = (
       {...restoreFocus}
     >
       <Command>
-        <CommandInput placeholder="Search actions..." className="h-6 m-1" />
+        <CommandInput
+          placeholder="Search actions..."
+          className="h-6 m-1"
+          {...getCellForDomProps(props.cellId)}
+        />
         <CommandList>
           <CommandEmpty>No results</CommandEmpty>
           {actions.map((group, i) => (

package/src/components/editor/controls/Controls.tsx

@@ -54,9 +54,7 @@ export const Controls = ({
   onRun,
   connectionState,
   running,
-  appConfig,
 }: ControlsProps): JSX.Element => {
-  const appWidth = appConfig.width;
   const undoAvailable = useAtomValue(canUndoDeletesAtom);
   const needsRun = useAtomValue(needsRunAtom);
   const { undoDeleteCell } = useCellActions();
@@ -103,12 +101,7 @@ export const Controls = ({
         </div>
       )}
 
-      <div
-        className={cn(
-          bottomRightControls,
-          appWidth === "compact" && "xl:flex-row items-end",
-        )}
-      >
+      <div className={cn(bottomRightControls)}>
         <HideInKioskMode>
           <SaveComponent kioskMode={false} />
         </HideInKioskMode>

package/src/components/editor/file-tree/file-explorer.tsx

@@ -1,6 +1,6 @@
 /* Copyright 2026 Marimo. All rights reserved. */
 
-import { useAtom } from "jotai";
+import { useAtom, useAtomValue } from "jotai";
 import { atomWithStorage } from "jotai/utils";
 import {
   ArrowLeftIcon,
@@ -47,6 +47,7 @@ import { Tooltip } from "@/components/ui/tooltip";
 import { toast } from "@/components/ui/use-toast";
 import { useCellActions } from "@/core/cells/cells";
 import { useLastFocusedCellId } from "@/core/cells/focus";
+import { disableFileDownloadsAtom } from "@/core/config/config";
 import { useRequestClient } from "@/core/network/requests";
 import type { FileInfo } from "@/core/network/types";
 import { isWasm } from "@/core/wasm/utils";
@@ -396,6 +397,7 @@ const Edit = ({ node }: { node: NodeApi<FileInfo> }) => {
 const Node = ({ node, style, dragHandle }: NodeRendererProps<FileInfo>) => {
   const { openFile, sendCreateFileOrFolder, sendFileDetails } =
     useRequestClient();
+  const disableFileDownloads = useAtomValue(disableFileDownloadsAtom);
 
   const fileType: FileType = node.data.isDirectory
     ? "directory"
@@ -611,7 +613,7 @@ const Node = ({ node, style, dragHandle }: NodeRendererProps<FileInfo>) => {
           </>
         )}
         <DropdownMenuSeparator />
-        {!node.data.isDirectory && (
+        {!node.data.isDirectory && !disableFileDownloads && (
           <>
             <DropdownMenuItem
               onSelect={async () => {

package/src/components/editor/file-tree/file-viewer.tsx

@@ -15,7 +15,7 @@ import { Suspense, useEffect, useRef, useState } from "react";
 import { renderShortcut } from "@/components/shortcuts/renderShortcut";
 import { Alert, AlertDescription } from "@/components/ui/alert";
 import { Tooltip } from "@/components/ui/tooltip";
-import { hotkeysAtom } from "@/core/config/config";
+import { disableFileDownloadsAtom, hotkeysAtom } from "@/core/config/config";
 import { useRequestClient } from "@/core/network/requests";
 import type { FileInfo } from "@/core/network/types";
 import { filenameAtom } from "@/core/saving/file-state";
@@ -49,6 +49,7 @@ export const FileViewer: React.FC<Props> = ({ file, onOpenNotebook }) => {
   const { theme } = useTheme();
   const { sendFileDetails, sendUpdateFile } = useRequestClient();
   const hotkeys = useAtomValue(hotkeysAtom);
+  const disableFileDownloads = useAtomValue(disableFileDownloadsAtom);
   const currentNotebookFilename = useAtomValue(filenameAtom);
   // undefined value means not modified yet
   const [internalValue, setInternalValue] = useState<string>("");
@@ -152,11 +153,13 @@ export const FileViewer: React.FC<Props> = ({ file, onOpenNotebook }) => {
           </Button>
         </Tooltip>
       )}
-      <Tooltip content="Download">
-        <Button size="small" onClick={handleDownload}>
-          <DownloadIcon />
-        </Button>
-      </Tooltip>
+      {!disableFileDownloads && (
+        <Tooltip content="Download">
+          <Button size="small" onClick={handleDownload}>
+            <DownloadIcon />
+          </Button>
+        </Tooltip>
+      )}
       {!isMedia(mimeType) && (
         <>
           <Tooltip content="Copy contents to clipboard">

package/src/components/editor/navigation/navigation.ts

@@ -10,6 +10,7 @@ import {
 import { useAtomValue, useSetAtom, useStore } from "jotai";
 import { useMemo } from "react";
 import { mergeProps, useFocusWithin, useKeyboard } from "react-aria";
+import { DATA_FOR_CELL_ID } from "@/components/data-table/cell-utils";
 import { aiCompletionCellAtom } from "@/core/ai/state";
 import { cellIdsAtom, notebookAtom, useCellActions } from "@/core/cells/cells";
 import { useCellFocusActions } from "@/core/cells/focus";
@@ -101,7 +102,23 @@ function useCellFocusProps(
       // On focus, set the last focused cell id.
       focusActions.focusCell({ cellId });
     },
-    onBlurWithin: () => {
+    onBlurWithin: (e) => {
+      // Check if blur is happening because of vim search panel interaction
+      if (isInVimPanel(e.relatedTarget) || isInVimPanel(e.target)) {
+        // Don't hide code if we're just interacting with vim search
+        return;
+      }
+
+      // If the related target is for a cell id (data-for-cell-id), then we don't want to hide the code, otherwise it might
+      // close the dropdown.
+      if (
+        getDataForCellId(e.relatedTarget) === cellId ||
+        getDataForCellId(e.relatedTarget?.closest(`[${DATA_FOR_CELL_ID}]`)) ===
+          cellId
+      ) {
+        return;
+      }
+
       // On blur, hide the code if it was temporarily shown.
       temporarilyShownCodeActions.remove(cellId);
       actions.markTouched({ cellId });
@@ -116,6 +133,27 @@ function useCellFocusProps(
   return focusWithinProps;
 }
 
+function isInVimPanel(element: Element | null): boolean {
+  if (!element) {
+    return false;
+  }
+  if (element instanceof HTMLElement) {
+    return element.closest(".cm-vim-panel") !== null;
+  }
+  return false;
+}
+
+function getDataForCellId(element: Element | null | undefined): CellId | null {
+  if (!element) {
+    return null;
+  }
+  const cellId = element.getAttribute(DATA_FOR_CELL_ID);
+  if (!cellId) {
+    return null;
+  }
+  return cellId as CellId;
+}
+
 type KeymapHandlers = Record<string, () => boolean>;
 
 /**

package/src/components/editor/renderMimeIcon.tsx

@@ -29,6 +29,8 @@ export function renderMimeIcon(mime: string) {
       return "📝";
     case "application/vnd.vegalite.v5+json":
     case "application/vnd.vega.v5+json":
+    case "application/vnd.vegalite.v6+json":
+    case "application/vnd.vega.v6+json":
       return "📊";
     case "application/vnd.marimo+mimebundle":
       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>();