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

package/src/components/app-config/ai-config.tsx

@@ -115,6 +115,7 @@ interface ApiKeyProps {
   placeholder: string;
   testId: string;
   description?: React.ReactNode;
+  onChange?: (value: string) => void;
 }
 
 export const ApiKey: React.FC<ApiKeyProps> = ({
@@ -124,6 +125,7 @@ export const ApiKey: React.FC<ApiKeyProps> = ({
   placeholder,
   testId,
   description,
+  onChange,
 }) => {
   return (
     <FormField
@@ -141,11 +143,12 @@ export const ApiKey: React.FC<ApiKeyProps> = ({
                 placeholder={placeholder}
                 type="password"
                 {...field}
-                value={asStringOrUndefined(field.value)}
+                value={asStringOrEmpty(field.value)}
                 onChange={(e) => {
                   const value = e.target.value;
                   if (!value.includes("*")) {
                     field.onChange(value);
+                    onChange?.(value);
                   }
                 }}
               />
@@ -168,12 +171,12 @@ interface BaseUrlProps {
   testId: string;
   description?: React.ReactNode;
   disabled?: boolean;
-  defaultValue?: string;
+  onChange?: (value: string) => void;
 }
 
-function asStringOrUndefined<T>(value: T): string | undefined {
+function asStringOrEmpty<T>(value: T): string {
   if (value == null) {
-    return undefined;
+    return "";
   }
 
   if (typeof value === "string") {
@@ -191,13 +194,12 @@ export const BaseUrl: React.FC<BaseUrlProps> = ({
   testId,
   description,
   disabled = false,
-  defaultValue,
+  onChange,
 }) => {
   return (
     <FormField
       control={form.control}
       name={name}
-      disabled={disabled}
       render={({ field }) => (
         <div className="flex flex-col space-y-1">
           <FormItem className={formItemClasses}>
@@ -208,9 +210,13 @@ export const BaseUrl: React.FC<BaseUrlProps> = ({
                 rootClassName="flex-1"
                 className="m-0 inline-flex h-7"
                 placeholder={placeholder}
-                defaultValue={defaultValue}
                 {...field}
-                value={asStringOrUndefined(field.value)}
+                value={asStringOrEmpty(field.value)}
+                disabled={disabled}
+                onChange={(e) => {
+                  field.onChange(e.target.value);
+                  onChange?.(e.target.value);
+                }}
               />
             </FormControl>
             <FormMessage />
@@ -252,9 +258,8 @@ export const ModelSelector: React.FC<ModelSelectorProps> = ({
     <FormField
       control={form.control}
       name={name}
-      disabled={disabled}
       render={({ field }) => {
-        const value = asStringOrUndefined(field.value);
+        const value = asStringOrEmpty(field.value);
 
         const selectModel = (modelId: QualifiedModelId) => {
           field.onChange(modelId);
@@ -286,7 +291,8 @@ export const ModelSelector: React.FC<ModelSelectorProps> = ({
                         className="w-full border-border shadow-none focus-visible:shadow-xs"
                         placeholder={placeholder}
                         {...field}
-                        value={asStringOrUndefined(field.value)}
+                        value={asStringOrEmpty(field.value)}
+                        disabled={disabled}
                         onKeyDown={Events.stopPropagation()}
                       />
                       {value && (
@@ -339,7 +345,6 @@ export const ProviderSelect: React.FC<ProviderSelectProps> = ({
     <FormField
       control={form.control}
       name={name}
-      disabled={disabled}
       render={({ field }) => (
         <div className="flex flex-col space-y-1">
           <FormItem className={formItemClasses}>
@@ -354,14 +359,14 @@ export const ProviderSelect: React.FC<ProviderSelectProps> = ({
                     field.onChange(e.target.value);
                   }
                 }}
-                value={asStringOrUndefined(
+                value={asStringOrEmpty(
                   field.value === true
                     ? "github"
                     : field.value === false
                       ? "none"
                       : field.value,
                 )}
-                disabled={field.disabled}
+                disabled={disabled}
                 className="inline-flex mr-2"
               >
                 {options.map((option) => (
@@ -715,6 +720,22 @@ export const CustomProvidersConfig: React.FC<AiConfigProps> = ({
           </div>
         );
 
+        // Update a provider field by updating the entire custom_providers object.
+        // As this config will be replaced, it needs to be sent in its entirety.
+        const updateProviderField = (opts: {
+          providerName: string;
+          fieldName: keyof CustomProviderConfig;
+          value: string;
+        }) => {
+          field.onChange({
+            ...customProviders,
+            [opts.providerName]: {
+              ...customProviders[opts.providerName],
+              [opts.fieldName]: opts.value || undefined,
+            },
+          });
+        };
+
         const renderAccordionItem = ({
           providerName,
           providerConfig,
@@ -744,6 +765,13 @@ export const CustomProvidersConfig: React.FC<AiConfigProps> = ({
                 }
                 placeholder="sk-..."
                 testId={`custom-provider-${providerName}-api-key`}
+                onChange={(value) =>
+                  updateProviderField({
+                    providerName,
+                    fieldName: "api_key",
+                    value,
+                  })
+                }
               />
               <BaseUrl
                 form={form}
@@ -753,6 +781,13 @@ export const CustomProvidersConfig: React.FC<AiConfigProps> = ({
                 }
                 placeholder="https://api.example.com/v1"
                 testId={`custom-provider-${providerName}-base-url`}
+                onChange={(value) =>
+                  updateProviderField({
+                    providerName,
+                    fieldName: "base_url",
+                    value,
+                  })
+                }
               />
               <Button
                 variant="destructive"
@@ -918,7 +953,6 @@ export const AiProvidersConfig: React.FC<AiConfigProps> = ({
             config={config}
             name="ai.ollama.base_url"
             placeholder="http://localhost:11434/v1"
-            defaultValue="http://localhost:11434/v1"
             testId="ollama-base-url-input"
           />
         </AccordionFormItem>
@@ -1038,7 +1072,6 @@ export const AiProvidersConfig: React.FC<AiConfigProps> = ({
             config={config}
             name="ai.azure.base_url"
             placeholder="https://<your-resource-name>.openai.azure.com/openai/deployments/<deployment-name>?api-version=<api-version>"
-            defaultValue="https://<your-resource-name>.openai.azure.com/openai/deployments/<deployment-name>?api-version=<api-version>"
             testId="ai-azure-base-url-input"
           />
         </AccordionFormItem>
@@ -1427,6 +1460,13 @@ export const AiModelDisplayConfig: React.FC<AiConfigProps> = ({
 
   const deleteModel = useEvent((modelId: QualifiedModelId) => {
     const newModels = customModels.filter((id) => id !== modelId);
+    // Remove from displayed models if it's in there
+    const newDisplayedModels = currentDisplayedModels.filter(
+      (id) => id !== modelId,
+    );
+    form.setValue("ai.models.displayed_models", newDisplayedModels, {
+      shouldDirty: true,
+    });
     form.setValue("ai.models.custom_models", newModels, {
       shouldDirty: true,
     });

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 }));

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],