@marimo-team/frontend 0.19.8-dev9 → 0.19.9-dev0

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

@@ -7,14 +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,
-  PaperclipIcon,
+  type LucideIcon,
+  MessageCircleIcon,
+  NotebookText,
   PlusIcon,
-  SendIcon,
   SettingsIcon,
-  SquareIcon,
 } from "lucide-react";
 import { memo, useEffect, useRef, useState } from "react";
 import useEvent from "react-use-event-hook";
@@ -29,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,
@@ -45,7 +45,6 @@ import {
 import { useCellActions } from "@/core/cells/cells";
 import { aiAtom, aiEnabledAtom } from "@/core/config/config";
 import { DEFAULT_AI_MODEL } from "@/core/config/config-schema";
-import { FeatureFlagged } from "@/core/config/feature-flag";
 import { useRequestClient } from "@/core/network/requests";
 import { useRuntimeManager } from "@/core/runtime/config";
 import { ErrorBanner } from "@/plugins/impl/common/error-banner";
@@ -61,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 {
@@ -74,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;
@@ -238,58 +233,72 @@ const ChatInputFooter: React.FC<ChatInputFooterProps> = memo(
       value: CopilotMode;
       label: string;
       subtitle: string;
+      Icon: LucideIcon;
     }[] = [
+      {
+        value: "manual",
+        label: "Manual",
+        subtitle: "Pure chat, no tool usage",
+        Icon: MessageCircleIcon,
+      },
       {
         value: "ask",
         label: "Ask",
         subtitle:
           "Use AI with access to read-only tools like documentation search",
-      },
-      {
-        value: "manual",
-        label: "Manual",
-        subtitle: "Pure chat, no tool usage",
+        Icon: NotebookText,
       },
       {
         value: "agent",
         label: "Agent (beta)",
         subtitle: "Use AI with access to read and write tools",
+        Icon: HatGlasses,
       },
     ];
 
     const isAttachmentSupported =
       PROVIDERS_THAT_SUPPORT_ATTACHMENTS.has(currentProvider);
 
+    const CurrentModeIcon = modeOptions.find(
+      (o) => o.value === currentMode,
+    )?.Icon;
+
     return (
       <TooltipProvider>
         <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}>
-                <SelectTrigger className="h-6 text-xs border-border shadow-none! ring-0! bg-muted hover:bg-muted/30 py-0 px-2 gap-1 capitalize">
-                  {currentMode}
-                </SelectTrigger>
-                <SelectContent>
-                  <SelectGroup>
-                    <SelectLabel>AI Mode</SelectLabel>
-                    {modeOptions.map((option) => (
-                      <SelectItem
-                        key={option.value}
-                        value={option.value}
-                        className="text-xs"
-                      >
-                        <div className="flex flex-col">
-                          {option.label}
-                          <div className="text-muted-foreground text-xs pt-1 block">
+            <Select value={currentMode} onValueChange={saveModeChange}>
+              <SelectTrigger className="h-6 text-xs border-border shadow-none! ring-0! bg-muted hover:bg-muted/30 py-0 px-2 gap-1.5">
+                {CurrentModeIcon && <CurrentModeIcon className="h-3 w-3" />}
+                <span className="capitalize">{currentMode}</span>
+              </SelectTrigger>
+              <SelectContent>
+                <SelectGroup>
+                  <SelectLabel className="text-xs uppercase tracking-wider text-muted-foreground/70 font-medium">
+                    AI Mode
+                  </SelectLabel>
+                  {modeOptions.map((option) => (
+                    <SelectItem
+                      key={option.value}
+                      value={option.value}
+                      className="text-xs py-1"
+                    >
+                      <div className="flex items-start gap-2.5">
+                        <span className="mt-1 text-muted-foreground">
+                          <option.Icon className="h-3 w-3" />
+                        </span>
+                        <div className="flex flex-col gap-0.5">
+                          <span className="font-semibold">{option.label}</span>
+                          <span className="text-muted-foreground">
                             {option.subtitle}
-                          </div>
+                          </span>
                         </div>
-                      </SelectItem>
-                    ))}
-                  </SelectGroup>
-                </SelectContent>
-              </Select>
-            </FeatureFlagged>
+                      </div>
+                    </SelectItem>
+                  ))}
+                </SelectGroup>
+              </SelectContent>
+            </Select>
             <AIModelDropdown
               placeholder="Model"
               triggerClassName="h-6 text-xs shadow-none! ring-0! bg-muted hover:bg-muted/30 rounded-sm"
@@ -299,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>
@@ -456,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);
@@ -545,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
@@ -631,7 +576,7 @@ const ChatPanelBody = () => {
         ...(fileParts ?? []),
       ],
     });
-    setFiles(undefined);
+    clearFiles();
     setInput("");
   };
 
@@ -639,7 +584,7 @@ const ChatPanelBody = () => {
     setActiveChat(null);
     setInput("");
     setNewThreadInput("");
-    setFiles(undefined);
+    clearFiles();
   });
 
   const handleMessageEdit = useEvent((index: number, newValue: string) => {
@@ -670,7 +615,7 @@ const ChatPanelBody = () => {
         files: fileParts,
       });
       setInput("");
-      setFiles(undefined);
+      clearFiles();
     },
   );
 
@@ -703,7 +648,7 @@ const ChatPanelBody = () => {
       isLoading={isLoading}
       onStop={stop}
       fileInputRef={fileInputRef}
-      onAddFiles={onAddFiles}
+      onAddFiles={addFiles}
       onClose={handleOnCloseThread}
     />
   ) : (
@@ -716,7 +661,7 @@ const ChatPanelBody = () => {
       onStop={stop}
       onClose={() => newMessageInputRef.current?.editor?.blur()}
       fileInputRef={fileInputRef}
-      onAddFiles={onAddFiles}
+      onAddFiles={addFiles}
     />
   );
 
@@ -777,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/actions/useCellActionButton.tsx

@@ -43,6 +43,7 @@ import type { CellData } from "@/core/cells/types";
 import { formatEditorViews } from "@/core/codemirror/format";
 import { toggleToLanguage } from "@/core/codemirror/language/commands";
 import { switchLanguage } from "@/core/codemirror/language/extension";
+import { MARKDOWN_INITIAL_HIDE_CODE } from "@/core/codemirror/language/languages/markdown";
 import {
   aiEnabledAtom,
   appWidthAtom,
@@ -85,6 +86,7 @@ export function useCellActionButtons({ cell, closePopover }: Props) {
     sendToBottom,
     addColumnBreakpoint,
     clearCellOutput,
+    markUntouched,
   } = useCellActions();
   const splitCell = useSplitCellCallback();
   const runCell = useRunCell(cell?.cellId);
@@ -209,7 +211,7 @@ export function useCellActionButtons({ cell, closePopover }: Props) {
         icon: <MarkdownIcon />,
         label: "Convert to Markdown",
         hotkey: "cell.viewAsMarkdown",
-        handle: () => {
+        handle: async () => {
           const editorView = getEditorView();
           if (!editorView) {
             return;
@@ -219,6 +221,17 @@ export function useCellActionButtons({ cell, closePopover }: Props) {
             language: "markdown",
             keepCodeAsIs: false,
           });
+          // Code stays visible until the user blurs the cell
+          if (!config.hide_code && MARKDOWN_INITIAL_HIDE_CODE) {
+            await saveCellConfig({
+              configs: { [cellId]: { hide_code: MARKDOWN_INITIAL_HIDE_CODE } },
+            });
+            updateCellConfig({
+              cellId,
+              config: { hide_code: MARKDOWN_INITIAL_HIDE_CODE },
+            });
+            markUntouched({ cellId });
+          }
         },
         hidden: isSetupCell,
       },

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 {