@marimo-team/frontend 0.19.8-dev41 → 0.19.8-dev49

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

@@ -1,9 +1,23 @@
 /* Copyright 2026 Marimo. All rights reserved. */
 
 import type { FileUIPart } from "ai";
-import { FileIcon, FileTextIcon, ImageIcon, XIcon } from "lucide-react";
+import {
+  AtSignIcon,
+  FileIcon,
+  FileTextIcon,
+  ImageIcon,
+  PaperclipIcon,
+  SendHorizontalIcon,
+  StopCircleIcon,
+  XIcon,
+} from "lucide-react";
 import { useState } from "react";
 import { cn } from "@/utils/cn";
+import { Spinner } from "../icons/spinner";
+import { Button } from "../ui/button";
+import { Input } from "../ui/input";
+import { Tooltip } from "../ui/tooltip";
+import { SUPPORTED_ATTACHMENT_TYPES } from "./chat-utils";
 
 export const AttachmentRenderer = ({
   attachment,
@@ -58,6 +72,105 @@ export const FileAttachmentPill = ({
   );
 };
 
+export const SendButton = ({
+  isLoading,
+  onStop,
+  onSendClick,
+  isEmpty,
+  showStopLabel = false, // Show a stop label and spinner instead of just the stop icon when loading
+}: {
+  isLoading: boolean;
+  onStop: () => void;
+  onSendClick: () => void;
+  isEmpty: boolean;
+  showStopLabel?: boolean;
+}) => {
+  const loadingContent = showStopLabel ? (
+    <div className="flex flex-row items-center gap-1 px-1.5">
+      <span className="text-xs text-error">Stop</span>
+      <Spinner size="small" />
+    </div>
+  ) : (
+    <StopCircleIcon className="size-4 text-error" />
+  );
+
+  return (
+    <Tooltip content={isLoading ? "Stop" : "Submit"}>
+      <Button
+        variant="text"
+        size="sm"
+        className="h-6 min-w-6 p-0 hover:bg-muted/30 cursor-pointer"
+        onClick={isLoading ? onStop : onSendClick}
+        disabled={isLoading ? false : isEmpty}
+      >
+        {isLoading ? (
+          loadingContent
+        ) : (
+          <SendHorizontalIcon className="h-3.5 w-3.5" />
+        )}
+      </Button>
+    </Tooltip>
+  );
+};
+
+export const AddContextButton = ({
+  handleAddContext,
+  isLoading,
+}: {
+  handleAddContext: () => void;
+  isLoading: boolean;
+}) => {
+  return (
+    <Tooltip content="Add context">
+      <Button
+        variant="text"
+        size="icon"
+        onClick={handleAddContext}
+        disabled={isLoading}
+      >
+        <AtSignIcon className="h-3.5 w-3.5" />
+      </Button>
+    </Tooltip>
+  );
+};
+
+export const AttachFileButton = ({
+  fileInputRef,
+  isLoading,
+  onAddFiles,
+}: {
+  fileInputRef: React.RefObject<HTMLInputElement | null>;
+  isLoading: boolean;
+  onAddFiles: (files: File[]) => void;
+}) => {
+  return (
+    <>
+      <Tooltip content="Attach a file">
+        <Button
+          variant="text"
+          size="icon"
+          onClick={() => fileInputRef.current?.click()}
+          disabled={isLoading}
+        >
+          <PaperclipIcon className="h-3.5 w-3.5" />
+        </Button>
+      </Tooltip>
+      <Input
+        ref={fileInputRef}
+        type="file"
+        multiple={true}
+        hidden={true}
+        onChange={(event: React.ChangeEvent<HTMLInputElement>) => {
+          if (event.target.files) {
+            onAddFiles([...event.target.files]);
+          }
+        }}
+        accept={SUPPORTED_ATTACHMENT_TYPES.join(",")}
+      />
+    </>
+  );
+};
+
 function renderFileIcon(file: File): React.ReactNode {
   const classNames = "h-3 w-3 mt-0.5";
 

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

@@ -7,18 +7,14 @@ import type { ReactCodeMirrorRef } from "@uiw/react-codemirror";
 import { DefaultChatTransport, type FileUIPart, type TextUIPart } from "ai";
 import { useAtom, useAtomValue, useSetAtom, useStore } from "jotai";
 import {
-  AtSignIcon,
   BotMessageSquareIcon,
   HatGlasses,
   Loader2,
   type LucideIcon,
   MessageCircleIcon,
   NotebookText,
-  PaperclipIcon,
   PlusIcon,
-  SendIcon,
   SettingsIcon,
-  SquareIcon,
 } from "lucide-react";
 import { memo, useEffect, useRef, useState } from "react";
 import useEvent from "react-use-event-hook";
@@ -33,7 +29,7 @@ import {
 } from "@/components/ui/select";
 import { replaceMessagesInChat } from "@/core/ai/chat-utils";
 import { useModelChange } from "@/core/ai/config";
-import { AiModelId, type ProviderId } from "@/core/ai/ids/ids";
+import { AiModelId } from "@/core/ai/ids/ids";
 import { useStagedAICellsActions } from "@/core/ai/staged-cells";
 import {
   activeChatAtom,
@@ -64,10 +60,14 @@ import {
 import { PanelEmptyState } from "../editor/chrome/panels/empty-state";
 import { CopyClipboardIcon } from "../icons/copy-icon";
 import { MCPStatusIndicator } from "../mcp/mcp-status-indicator";
-import { Input } from "../ui/input";
 import { Tooltip, TooltipProvider } from "../ui/tooltip";
-import { toast } from "../ui/use-toast";
-import { AttachmentRenderer, FileAttachmentPill } from "./chat-components";
+import {
+  AddContextButton,
+  AttachFileButton,
+  AttachmentRenderer,
+  FileAttachmentPill,
+  SendButton,
+} from "./chat-components";
 import { renderUIMessage } from "./chat-display";
 import { ChatHistoryPopover } from "./chat-history-popover";
 import {
@@ -77,21 +77,13 @@ import {
   handleToolCall,
   hasPendingToolCalls,
   isLastMessageReasoning,
+  PROVIDERS_THAT_SUPPORT_ATTACHMENTS,
+  useFileState,
 } from "./chat-utils";
 
 // Default mode for the AI
 const DEFAULT_MODE = "manual";
 
-// We need to modify the backend to support attachments for other providers
-// And other types
-const PROVIDERS_THAT_SUPPORT_ATTACHMENTS = new Set<ProviderId>([
-  "openai",
-  "google",
-  "anthropic",
-]);
-const SUPPORTED_ATTACHMENT_TYPES = ["image/*", "text/*"];
-const MAX_ATTACHMENT_SIZE = 1024 * 1024 * 50; // 50MB
-
 interface ChatHeaderProps {
   onNewChat: () => void;
   activeChatId: ChatId | undefined;
@@ -316,60 +308,23 @@ const ChatInputFooter: React.FC<ChatInputFooterProps> = memo(
             />
           </div>
           <div className="flex flex-row">
-            <Tooltip content="Add context">
-              <Button
-                variant="text"
-                size="icon"
-                onClick={onAddContext}
-                disabled={isLoading}
-              >
-                <AtSignIcon className="h-3.5 w-3.5" />
-              </Button>
-            </Tooltip>
+            <AddContextButton
+              handleAddContext={onAddContext}
+              isLoading={isLoading}
+            />
             {isAttachmentSupported && (
-              <>
-                <Tooltip content="Attach a file">
-                  <Button
-                    variant="text"
-                    size="icon"
-                    className="cursor-pointer"
-                    onClick={() => fileInputRef.current?.click()}
-                    title="Attach a file"
-                    disabled={isLoading}
-                  >
-                    <PaperclipIcon className="h-3.5 w-3.5" />
-                  </Button>
-                </Tooltip>
-                <Input
-                  ref={fileInputRef}
-                  type="file"
-                  multiple={true}
-                  hidden={true}
-                  onChange={(event) => {
-                    if (event.target.files) {
-                      onAddFiles([...event.target.files]);
-                    }
-                  }}
-                  accept={SUPPORTED_ATTACHMENT_TYPES.join(",")}
-                />
-              </>
+              <AttachFileButton
+                fileInputRef={fileInputRef}
+                isLoading={isLoading}
+                onAddFiles={onAddFiles}
+              />
             )}
-
-            <Tooltip content={isLoading ? "Stop" : "Submit"}>
-              <Button
-                variant="text"
-                size="sm"
-                className="h-6 w-6 p-0 hover:bg-muted/30 cursor-pointer"
-                onClick={isLoading ? onStop : onSendClick}
-                disabled={isLoading ? false : isEmpty}
-              >
-                {isLoading ? (
-                  <SquareIcon className="h-3 w-3 fill-current" />
-                ) : (
-                  <SendIcon className="h-3 w-3" />
-                )}
-              </Button>
-            </Tooltip>
+            <SendButton
+              isLoading={isLoading}
+              onStop={onStop}
+              onSendClick={onSendClick}
+              isEmpty={isEmpty}
+            />
           </div>
         </div>
       </TooltipProvider>
@@ -473,7 +428,7 @@ const ChatPanelBody = () => {
   const [activeChat, setActiveChat] = useAtom(activeChatAtom);
   const [input, setInput] = useState("");
   const [newThreadInput, setNewThreadInput] = useState("");
-  const [files, setFiles] = useState<File[]>();
+  const { files, addFiles, clearFiles, removeFile } = useFileState();
   const newThreadInputRef = useRef<ReactCodeMirrorRef>(null);
   const newMessageInputRef = useRef<ReactCodeMirrorRef>(null);
   const scrollContainerRef = useRef<HTMLDivElement>(null);
@@ -562,33 +517,6 @@ const ChatPanelBody = () => {
     },
   });
 
-  const onAddFiles = useEvent((files: File[]) => {
-    if (files.length === 0) {
-      return;
-    }
-
-    let fileSize = 0;
-    for (const file of files) {
-      fileSize += file.size;
-    }
-
-    if (fileSize > MAX_ATTACHMENT_SIZE) {
-      toast({
-        title: "File size exceeds 50MB limit",
-        description: "Please remove some files and try again.",
-      });
-      return;
-    }
-
-    setFiles((prev) => [...(prev ?? []), ...files]);
-  });
-
-  const removeFile = useEvent((file: File) => {
-    if (files) {
-      setFiles(files.filter((f) => f !== file));
-    }
-  });
-
   const isLoading = status === "submitted" || status === "streaming";
 
   // Check if we're currently streaming reasoning in the latest message
@@ -648,7 +576,7 @@ const ChatPanelBody = () => {
         ...(fileParts ?? []),
       ],
     });
-    setFiles(undefined);
+    clearFiles();
     setInput("");
   };
 
@@ -656,7 +584,7 @@ const ChatPanelBody = () => {
     setActiveChat(null);
     setInput("");
     setNewThreadInput("");
-    setFiles(undefined);
+    clearFiles();
   });
 
   const handleMessageEdit = useEvent((index: number, newValue: string) => {
@@ -687,7 +615,7 @@ const ChatPanelBody = () => {
         files: fileParts,
       });
       setInput("");
-      setFiles(undefined);
+      clearFiles();
     },
   );
 
@@ -720,7 +648,7 @@ const ChatPanelBody = () => {
       isLoading={isLoading}
       onStop={stop}
       fileInputRef={fileInputRef}
-      onAddFiles={onAddFiles}
+      onAddFiles={addFiles}
       onClose={handleOnCloseThread}
     />
   ) : (
@@ -733,7 +661,7 @@ const ChatPanelBody = () => {
       onStop={stop}
       onClose={() => newMessageInputRef.current?.editor?.blur()}
       fileInputRef={fileInputRef}
-      onAddFiles={onAddFiles}
+      onAddFiles={addFiles}
     />
   );
 
@@ -794,7 +722,7 @@ const ChatPanelBody = () => {
 
         {error && (
           <div className="flex items-center justify-center space-x-2 mb-4">
-            <ErrorBanner error={error} />
+            <ErrorBanner error={error || new Error("Unknown error")} />
             <Button variant="outline" size="sm" onClick={handleReload}>
               Retry
             </Button>

package/src/components/chat/chat-utils.ts

@@ -2,6 +2,9 @@
 
 import type { components } from "@marimo-team/marimo-api";
 import type { FileUIPart, ToolUIPart, UIMessage } from "ai";
+import { useState } from "react";
+import useEvent from "react-use-event-hook";
+import type { ProviderId } from "@/core/ai/ids/ids";
 import type { ToolNotebookContext } from "@/core/ai/tools/base";
 import { FRONTEND_TOOL_REGISTRY } from "@/core/ai/tools/registry";
 import type {
@@ -11,6 +14,17 @@ import type {
 import { blobToString } from "@/utils/fileToBase64";
 import { Logger } from "@/utils/Logger";
 import { getAICompletionBodyWithAttachments } from "../editor/ai/completion-utils";
+import { toast } from "../ui/use-toast";
+
+// We need to modify the backend to support attachments for other providers
+// And other types
+export const PROVIDERS_THAT_SUPPORT_ATTACHMENTS = new Set<ProviderId>([
+  "openai",
+  "google",
+  "anthropic",
+]);
+export const SUPPORTED_ATTACHMENT_TYPES = ["image/*", "text/*"];
+const MAX_ATTACHMENT_SIZE = 1024 * 1024 * 50; // 50MB
 
 export function generateChatTitle(message: string): string {
   return message.length > 50 ? `${message.slice(0, 50)}...` : message;
@@ -198,3 +212,31 @@ export function hasPendingToolCalls(messages: UIMessage[]): boolean {
   // Only auto-send if we have completed tool calls and there is no reply yet
   return allToolCallsCompleted && !hasTextContent;
 }
+
+export function useFileState() {
+  const [files, setFiles] = useState<File[]>([]);
+
+  const addFiles = useEvent((newFiles: File[]) => {
+    if (newFiles.length === 0) {
+      return;
+    }
+
+    const totalSize = newFiles.reduce((size, file) => size + file.size, 0);
+    if (totalSize > MAX_ATTACHMENT_SIZE) {
+      toast({
+        title: "File size exceeded",
+        description: "Attachments must be under 50 MB",
+        variant: "danger",
+      });
+      return;
+    }
+
+    setFiles((prev) => [...prev, ...newFiles]);
+  });
+
+  const clearFiles = () => setFiles([]);
+  const removeFile = (fileToRemove: File) =>
+    setFiles((prev) => prev.filter((f) => f !== fileToRemove));
+
+  return { files, addFiles, clearFiles, removeFile };
+}

package/src/components/editor/ai/add-cell-with-ai.tsx

@@ -21,8 +21,6 @@ import { atomWithStorage } from "jotai/utils";
 import {
   ChevronsUpDown,
   DatabaseIcon,
-  Loader2Icon,
-  SendHorizontal,
   SparklesIcon,
   XIcon,
 } from "lucide-react";
@@ -30,9 +28,18 @@ import { useMemo, useRef, useState } from "react";
 import useEvent from "react-use-event-hook";
 import { z } from "zod";
 import { AIModelDropdown } from "@/components/ai/ai-model-dropdown";
+import {
+  AddContextButton,
+  AttachFileButton,
+  FileAttachmentPill,
+  SendButton,
+} from "@/components/chat/chat-components";
 import {
   buildCompletionRequestBody,
+  convertToFileUIPart,
   handleToolCall,
+  PROVIDERS_THAT_SUPPORT_ATTACHMENTS,
+  useFileState,
 } from "@/components/chat/chat-utils";
 import { Button } from "@/components/ui/button";
 import {
@@ -43,10 +50,13 @@ import {
   DropdownMenuTrigger,
 } from "@/components/ui/dropdown-menu";
 import { toast } from "@/components/ui/use-toast";
+import { AiModelId } from "@/core/ai/ids/ids";
 import { stagedAICellsAtom, useStagedCells } from "@/core/ai/staged-cells";
 import type { ToolNotebookContext } from "@/core/ai/tools/base";
 import { useCellActions } from "@/core/cells/cells";
 import { resourceExtension } from "@/core/codemirror/ai/resources";
+import { aiAtom } from "@/core/config/config";
+import { DEFAULT_AI_MODEL } from "@/core/config/config-schema";
 import { useRequestClient } from "@/core/network/requests";
 import type { AiCompletionRequest } from "@/core/network/types";
 import { useRuntimeManager } from "@/core/runtime/config";
@@ -57,7 +67,11 @@ import { jotaiJsonStorage } from "@/utils/storage/jotai";
 import { ZodLocalStorage } from "@/utils/storage/typed";
 import { PythonIcon } from "../cell/code/icons";
 import { createAiCompletionOnKeydown } from "./completion-handlers";
-import { CONTEXT_TRIGGER, mentionsCompletionSource } from "./completion-utils";
+import {
+  addContextCompletion,
+  CONTEXT_TRIGGER,
+  mentionsCompletionSource,
+} from "./completion-utils";
 import { StreamingChunkTransport } from "./transport/chat-transport";
 
 // Persist across sessions
@@ -89,6 +103,10 @@ export const AddCellWithAI: React.FC<{
   const stagedAICells = useAtomValue(stagedAICellsAtom);
   const inputRef = useRef<ReactCodeMirrorRef>(null);
 
+  const fileInputRef = useRef<HTMLInputElement>(null);
+  const { files, addFiles, removeFile } = useFileState();
+  const aiConfig = useAtomValue(aiAtom);
+
   const { createNewCell, prepareForRun } = useCellActions();
   const toolContext: ToolNotebookContext = {
     store,
@@ -149,14 +167,21 @@ export const AddCellWithAI: React.FC<{
   const isLoading = status === "streaming" || status === "submitted";
   const hasCompletion = stagedAICells.size > 0;
 
-  const submit = () => {
+  const currentModel = aiConfig?.models?.edit_model || DEFAULT_AI_MODEL;
+  const currentProvider = AiModelId.parse(currentModel).providerId;
+  const isAttachmentSupported =
+    PROVIDERS_THAT_SUPPORT_ATTACHMENTS.has(currentProvider);
+
+  const submit = async () => {
     if (!isLoading) {
       if (inputRef.current?.view) {
         storePrompt(inputRef.current.view);
       }
       // TODO: When we have conversations, don't delete existing cells
       deleteAllStagedCells();
-      sendMessage({ text: input });
+
+      const fileParts = files ? await convertToFileUIPart(files) : undefined;
+      sendMessage({ text: input, files: fileParts });
     }
   };
 
@@ -176,19 +201,17 @@ export const AddCellWithAI: React.FC<{
 
   const languageDropdown = (
     <DropdownMenu modal={false}>
-      <DropdownMenuTrigger asChild={true}>
-        <Button
-          variant="text"
-          className="ml-2"
-          size="xs"
-          data-testid="language-button"
-        >
-          {language === "python" ? pythonIcon : sqlIcon}
-          <ChevronsUpDown className="ml-1 h-3.5 w-3.5 text-muted-foreground/70" />
-        </Button>
+      <DropdownMenuTrigger
+        className="flex items-center justify-between h-7 text-xs px-2 py-0.5 border rounded-md hover:text-accent-foreground"
+        data-testid="language-button"
+      >
+        {language === "python" ? pythonIcon : sqlIcon}
+        <ChevronsUpDown className="ml-1 h-3.5 w-3.5 text-muted-foreground/70" />
       </DropdownMenuTrigger>
       <DropdownMenuContent align="center">
-        <div className="px-2 py-1 font-semibold">Select language</div>
+        <div className="px-2 py-1 text-sm text-muted-foreground">
+          Select language
+        </div>
         <DropdownMenuSeparator />
         <DropdownMenuItem onClick={() => setLanguage("python")}>
           {pythonIcon}
@@ -230,54 +253,63 @@ export const AddCellWithAI: React.FC<{
           hasCompletion,
         })}
       />
-      {isLoading && (
-        <Button
-          data-testid="stop-completion-button"
-          variant="text"
-          size="sm"
-          className="mb-0"
-          onClick={stop}
-        >
-          <Loader2Icon className="animate-spin mr-1" size={14} />
-          Stop
-        </Button>
-      )}
-      <Button variant="text" size="sm" onClick={submit} title="Submit">
-        <SendHorizontal className="size-4" />
-      </Button>
       <Button variant="text" size="sm" className="mb-0 px-1" onClick={onClose}>
         <XIcon className="size-4" />
       </Button>
     </div>
   );
 
-  return (
-    <div className={cn("flex flex-col w-full py-2")}>
-      {inputComponent}
-      <div className="flex flex-row justify-between -mt-1 ml-1 mr-3">
-        {!hasCompletion && (
-          <span className="text-xs text-muted-foreground px-3 flex flex-col gap-1 mt-2">
-            <span>
-              You can mention{" "}
-              <span className="text-(--cyan-11)">@dataframe</span> or{" "}
-              <span className="text-(--cyan-11)">@sql_table</span> to pull
-              additional context such as column names. Code from other cells is
-              automatically included.
-            </span>
-          </span>
+  const footerComponent = (
+    <div className="px-3 pt-1 flex flex-row items-center justify-between">
+      <div className="flex items-center gap-2">
+        <AIModelDropdown
+          triggerClassName="h-7 text-xs max-w-64"
+          iconSize="small"
+          forRole="edit"
+          showAddCustomModelDocs={true}
+        />
+        {languageDropdown}
+      </div>
+      <div className="flex flex-row items-center">
+        {files.length > 0 && (
+          <div className="flex flex-row gap-1 flex-wrap pr-1">
+            {files?.map((file, index) => (
+              <FileAttachmentPill
+                file={file}
+                key={`${file.name}-${index}`}
+                onRemove={() => removeFile(file)}
+              />
+            ))}
+          </div>
         )}
-        <div className="ml-auto flex items-center gap-1">
-          {languageDropdown}
-          <AIModelDropdown
-            triggerClassName="h-7 text-xs max-w-64"
-            iconSize="small"
-            forRole="edit"
-            showAddCustomModelDocs={true}
+        <AddContextButton
+          handleAddContext={() => addContextCompletion(inputRef)}
+          isLoading={isLoading}
+        />
+        {isAttachmentSupported && (
+          <AttachFileButton
+            fileInputRef={fileInputRef}
+            isLoading={isLoading}
+            onAddFiles={addFiles}
           />
-        </div>
+        )}
+        <SendButton
+          isLoading={isLoading}
+          onStop={stop}
+          onSendClick={submit}
+          isEmpty={!input.trim()}
+          showStopLabel={true}
+        />
       </div>
     </div>
   );
+
+  return (
+    <div className={cn("flex flex-col w-full py-2")}>
+      {inputComponent}
+      {footerComponent}
+    </div>
+  );
 };
 
 export interface AdditionalCompletions {