@marimo-team/islands 0.15.2 → 0.15.4

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

@@ -1,28 +1,23 @@
 /* Copyright 2024 Marimo. All rights reserved. */
 
+import type { UIMessage } from "@ai-sdk/react";
 import { useChat } from "@ai-sdk/react";
 import { storePrompt } from "@marimo-team/codemirror-ai";
 import type { ReactCodeMirrorRef } from "@uiw/react-codemirror";
-import type { Message } from "ai/react";
-import { useAtom, useAtomValue } from "jotai";
+import { DefaultChatTransport, type ToolUIPart } from "ai";
+import { useAtom, useAtomValue, useSetAtom } from "jotai";
 import {
+  AtSignIcon,
   BotMessageSquareIcon,
   ClockIcon,
   Loader2,
+  PaperclipIcon,
   PlusIcon,
   SendIcon,
   SettingsIcon,
   SquareIcon,
 } from "lucide-react";
-import {
-  type Dispatch,
-  memo,
-  type SetStateAction,
-  useEffect,
-  useMemo,
-  useRef,
-  useState,
-} from "react";
+import { memo, useEffect, useMemo, useRef, useState } from "react";
 import useEvent from "react-use-event-hook";
 import { Button } from "@/components/ui/button";
 import {
@@ -39,51 +34,79 @@ import {
   SelectLabel,
   SelectTrigger,
 } from "@/components/ui/select";
-import { addMessageToChat } from "@/core/ai/chat-utils";
+import { replaceMessagesInChat } from "@/core/ai/chat-utils";
 import { useModelChange } from "@/core/ai/config";
+import { AiModelId, type ProviderId } from "@/core/ai/ids/ids";
 import {
   activeChatAtom,
   type Chat,
   type ChatId,
-  type ChatState,
   chatStateAtom,
 } from "@/core/ai/state";
-import { getCodes } from "@/core/codemirror/copilot/getCodes";
 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 type { ChatMessage } from "@/plugins/impl/chat/types";
 import { ErrorBanner } from "@/plugins/impl/common/error-banner";
 import { cn } from "@/utils/cn";
 import { timeAgo } from "@/utils/dates";
 import { Logger } from "@/utils/Logger";
-import { generateUUID } from "@/utils/uuid";
 import { AIModelDropdown } from "../ai/ai-model-dropdown";
 import { useOpenSettingsToTab } from "../app-config/state";
 import { PromptInput } from "../editor/ai/add-cell-with-ai";
-import { getAICompletionBody } from "../editor/ai/completion-utils";
+import {
+  addContextCompletion,
+  CONTEXT_TRIGGER,
+  getAICompletionBodyWithAttachments,
+} from "../editor/ai/completion-utils";
 import { PanelEmptyState } from "../editor/chrome/panels/empty-state";
 import { CopyClipboardIcon } from "../icons/copy-icon";
+import { Input } from "../ui/input";
 import { Tooltip, TooltipProvider } from "../ui/tooltip";
+import { toast } from "../ui/use-toast";
+import { AttachmentRenderer, FileAttachmentPill } from "./chat-components";
+import {
+  convertToFileUIPart,
+  generateChatTitle,
+  isLastMessageReasoning,
+} from "./chat-utils";
 import { MarkdownRenderer } from "./markdown-renderer";
 import { ReasoningAccordion } from "./reasoning-accordion";
 import { ToolCallAccordion } from "./tool-call-accordion";
 
+// 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;
   setActiveChat: (id: ChatId | null) => void;
-  chats: Chat[];
 }
 
 const ChatHeader: React.FC<ChatHeaderProps> = ({
   onNewChat,
   activeChatId,
   setActiveChat,
-  chats,
 }) => {
   const { handleClick } = useOpenSettingsToTab();
+  const chatState = useAtomValue(chatStateAtom);
+  const chats = useMemo(() => {
+    return [...chatState.chats.values()].sort(
+      (a, b) => b.updatedAt - a.updatedAt,
+    );
+  }, [chatState.chats]);
 
   return (
     <div className="flex border-b px-2 py-1 justify-between shrink-0 items-center">
@@ -149,28 +172,32 @@ const ChatHeader: React.FC<ChatHeaderProps> = ({
 };
 
 interface ChatMessageProps {
-  message: Message;
+  message: UIMessage;
   index: number;
   onEdit: (index: number, newValue: string) => void;
-  setChatState: Dispatch<SetStateAction<ChatState>>;
-  chatState: ChatState;
   isStreamingReasoning: boolean;
   isLast: boolean;
 }
 
-const ChatMessage: React.FC<ChatMessageProps> = memo(
-  ({ message, index, onEdit, isStreamingReasoning, isLast }) => (
-    <div
-      className={cn(
-        "flex group relative",
-        message.role === "user" ? "justify-end" : "justify-start",
-      )}
-    >
-      {message.role === "user" ? (
+function isToolPart(part: UIMessage["parts"][number]): part is ToolUIPart {
+  return part.type.startsWith("tool-");
+}
+
+const ChatMessageDisplay: React.FC<ChatMessageProps> = memo(
+  ({ message, index, onEdit, isStreamingReasoning, isLast }) => {
+    const renderUserMessage = (message: UIMessage) => {
+      const textParts = message.parts?.filter((p) => p.type === "text");
+      const content = textParts?.map((p) => p.text).join("\n");
+      const fileParts = message.parts?.filter((p) => p.type === "file");
+
+      return (
         <div className="w-[95%] bg-background border p-1 rounded-sm">
+          {fileParts?.map((filePart, idx) => (
+            <AttachmentRenderer attachment={filePart} key={idx} />
+          ))}
           <PromptInput
             key={message.id}
-            value={message.content}
+            value={content}
             placeholder="Type your message..."
             onChange={() => {
               // noop
@@ -186,12 +213,31 @@ const ChatMessage: React.FC<ChatMessageProps> = memo(
             }}
           />
         </div>
-      ) : (
+      );
+    };
+
+    const renderOtherMessage = (message: UIMessage) => {
+      const textParts = message.parts.filter((p) => p.type === "text");
+      const content = textParts.map((p) => p.text).join("\n");
+
+      return (
         <div className="w-[95%] break-words">
           <div className="absolute right-1 top-1 opacity-0 group-hover:opacity-100 transition-opacity">
-            <CopyClipboardIcon className="h-3 w-3" value={message.content} />
+            <CopyClipboardIcon className="h-3 w-3" value={content || ""} />
           </div>
-          {message.parts?.map((part, i) => {
+          {message.parts.map((part, i) => {
+            if (isToolPart(part)) {
+              return (
+                <ToolCallAccordion
+                  key={i}
+                  index={i}
+                  toolName={part.type}
+                  result={part.output}
+                  state={part.state}
+                />
+              );
+            }
+
             switch (part.type) {
               case "text":
                 return <MarkdownRenderer key={i} content={part.text} />;
@@ -199,60 +245,100 @@ const ChatMessage: React.FC<ChatMessageProps> = memo(
               case "reasoning":
                 return (
                   <ReasoningAccordion
-                    reasoning={part.reasoning}
+                    reasoning={part.text}
                     key={i}
                     index={i}
                     isStreaming={
                       isLast &&
                       isStreamingReasoning &&
                       // If there are multiple reasoning parts, only show the last one
-                      i === (message.parts?.length || 0) - 1
+                      i === (message.parts.length || 0) - 1
                     }
                   />
                 );
 
-              case "tool-invocation":
+              case "dynamic-tool":
                 return (
                   <ToolCallAccordion
                     key={i}
                     index={i}
-                    toolName={part.toolInvocation.toolName}
-                    result={
-                      part.toolInvocation.state === "result"
-                        ? part.toolInvocation.result
-                        : null
-                    }
-                    state={part.toolInvocation.state}
+                    toolName={part.type}
+                    result={part.output}
+                    state={part.state}
                   />
                 );
 
+              // These are cryptographic signatures, so we don't need to render them
+              case "data-reasoning-signature":
+                return null;
+
               /* handle other part types … */
               default:
-                return null;
+                if (part.type.startsWith("data-")) {
+                  Logger.log("Found data part", part);
+                  return null;
+                }
+
+                Logger.error("Unhandled part type:", part.type);
+                try {
+                  return (
+                    <MarkdownRenderer
+                      key={i}
+                      content={JSON.stringify(part, null, 2)}
+                    />
+                  );
+                } catch (error) {
+                  Logger.error("Error rendering part:", part.type, error);
+                  return null;
+                }
             }
           })}
         </div>
-      )}
-    </div>
-  ),
+      );
+    };
+
+    return (
+      <div
+        className={cn(
+          "flex group relative",
+          message.role === "user" ? "justify-end" : "justify-start",
+        )}
+      >
+        {message.role === "user"
+          ? renderUserMessage(message)
+          : renderOtherMessage(message)}
+      </div>
+    );
+  },
 );
-ChatMessage.displayName = "ChatMessage";
+ChatMessageDisplay.displayName = "ChatMessage";
 
 interface ChatInputFooterProps {
   isEmpty: boolean;
   onSendClick: () => void;
   isLoading: boolean;
   onStop: () => void;
+  onAddFiles: (files: File[]) => void;
+  fileInputRef: React.RefObject<HTMLInputElement | null>;
+  onAddContext: () => void;
 }
 
-const DEFAULT_MODE = "manual";
-
 const ChatInputFooter: React.FC<ChatInputFooterProps> = memo(
-  ({ isEmpty, onSendClick, isLoading, onStop }) => {
+  ({
+    isEmpty,
+    onSendClick,
+    isLoading,
+    onStop,
+    fileInputRef,
+    onAddFiles,
+    onAddContext,
+  }) => {
     const ai = useAtomValue(aiAtom);
     const currentMode = ai?.mode || DEFAULT_MODE;
     const currentModel = ai?.models?.chat_model || DEFAULT_AI_MODEL;
-    const { saveModeChange, saveModelChange } = useModelChange();
+    const currentProvider = AiModelId.parse(currentModel).providerId;
+
+    const { saveModeChange } = useModelChange();
 
     const modeOptions = [
       {
@@ -268,59 +354,99 @@ const ChatInputFooter: React.FC<ChatInputFooterProps> = memo(
       },
     ];
 
+    const isAttachmentSupported =
+      PROVIDERS_THAT_SUPPORT_ATTACHMENTS.has(currentProvider);
+
     return (
-      <div className="px-3 py-2 border-t border-border/20 flex items-center justify-between">
-        <div className="flex items-center gap-2">
-          <FeatureFlagged feature="mcp_docs">
-            <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">
-                          {option.subtitle}
+      <TooltipProvider>
+        <div className="px-3 py-2 border-t border-border/20 flex flex-row items-center justify-between">
+          <div className="flex items-center gap-2">
+            <FeatureFlagged feature="mcp_docs">
+              <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">
+                            {option.subtitle}
+                          </div>
                         </div>
-                      </div>
-                    </SelectItem>
-                  ))}
-                </SelectGroup>
-              </SelectContent>
-            </Select>
-          </FeatureFlagged>
-          <AIModelDropdown
-            value={currentModel}
-            placeholder="Model"
-            onSelect={(model) => saveModelChange(model, "chat")}
-            triggerClassName="h-6 text-xs shadow-none! ring-0! bg-muted hover:bg-muted/30 rounded-sm"
-            iconSize="small"
-            showAddCustomModelDocs={true}
-            forRole="chat"
-          />
+                      </SelectItem>
+                    ))}
+                  </SelectGroup>
+                </SelectContent>
+              </Select>
+            </FeatureFlagged>
+            <AIModelDropdown
+              placeholder="Model"
+              triggerClassName="h-6 text-xs shadow-none! ring-0! bg-muted hover:bg-muted/30 rounded-sm"
+              iconSize="small"
+              showAddCustomModelDocs={true}
+              forRole="chat"
+            />
+          </div>
+          <div className="flex flex-row">
+            <Tooltip content="Add context">
+              <Button variant="text" size="icon" onClick={onAddContext}>
+                <AtSignIcon className="h-3.5 w-3.5" />
+              </Button>
+            </Tooltip>
+            {isAttachmentSupported && (
+              <>
+                <Tooltip content="Attach a file">
+                  <Button
+                    variant="text"
+                    size="icon"
+                    className="cursor-pointer"
+                    onClick={() => fileInputRef.current?.click()}
+                    title="Attach a file"
+                  >
+                    <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(",")}
+                />
+              </>
+            )}
+
+            <Tooltip content="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>
+          </div>
         </div>
-        <Button
-          variant="ghost"
-          size="sm"
-          className="h-6 w-6 p-0 hover:bg-muted/30"
-          onClick={isLoading ? onStop : onSendClick}
-          disabled={isLoading ? false : isEmpty}
-        >
-          {isLoading ? (
-            <SquareIcon className="h-3 w-3 fill-current" />
-          ) : (
-            <SendIcon className="h-3 w-3" />
-          )}
-        </Button>
-      </div>
+      </TooltipProvider>
     );
   },
 );
@@ -328,16 +454,33 @@ const ChatInputFooter: React.FC<ChatInputFooterProps> = memo(
 ChatInputFooter.displayName = "ChatInputFooter";
 
 interface ChatInputProps {
+  placeholder?: string;
   input: string;
+  inputClassName?: string;
   setInput: (value: string) => void;
   onSubmit: (e: KeyboardEvent | undefined, value: string) => void;
   inputRef: React.RefObject<ReactCodeMirrorRef | null>;
   isLoading: boolean;
   onStop: () => void;
+  onClose: () => void;
+  fileInputRef: React.RefObject<HTMLInputElement | null>;
+  onAddFiles: (files: File[]) => void;
 }
 
 const ChatInput: React.FC<ChatInputProps> = memo(
-  ({ input, setInput, onSubmit, inputRef, isLoading, onStop }) => {
+  ({
+    placeholder,
+    input,
+    inputClassName,
+    setInput,
+    onSubmit,
+    inputRef,
+    isLoading,
+    onStop,
+    fileInputRef,
+    onAddFiles,
+    onClose,
+  }) => {
     const handleSendClick = useEvent(() => {
       if (input.trim()) {
         onSubmit(undefined, input);
@@ -345,22 +488,26 @@ const ChatInput: React.FC<ChatInputProps> = memo(
     });
 
     return (
-      <div className="border-t relative shrink-0 min-h-[80px] flex flex-col">
-        <div className="px-2 py-3 flex-1">
+      <div className="relative shrink-0 min-h-[80px] flex flex-col border-t">
+        <div className={cn("px-2 py-3 flex-1", inputClassName)}>
           <PromptInput
             inputRef={inputRef}
             value={input}
             onChange={setInput}
             onSubmit={onSubmit}
-            onClose={() => inputRef.current?.editor?.blur()}
-            placeholder="Type your message..."
+            onClose={onClose}
+            onAddFiles={onAddFiles}
+            placeholder={placeholder || "Type your message..."}
           />
         </div>
         <ChatInputFooter
           isEmpty={!input.trim()}
+          onAddContext={() => addContextCompletion(inputRef)}
           onSendClick={handleSendClick}
           isLoading={isLoading}
           onStop={onStop}
+          fileInputRef={fileInputRef}
+          onAddFiles={onAddFiles}
         />
       </div>
     );
@@ -369,7 +516,7 @@ const ChatInput: React.FC<ChatInputProps> = memo(
 
 ChatInput.displayName = "ChatInput";
 
-export const ChatPanel = () => {
+const ChatPanel = () => {
   const aiEnabled = useAtomValue(aiEnabledAtom);
   const { handleClick } = useOpenSettingsToTab();
 
@@ -392,130 +539,172 @@ export const ChatPanel = () => {
 };
 
 const ChatPanelBody = () => {
-  const [chatState, setChatState] = useAtom(chatStateAtom);
+  const setChatState = useSetAtom(chatStateAtom);
   const [activeChat, setActiveChat] = useAtom(activeChatAtom);
+  const [input, setInput] = useState("");
   const [newThreadInput, setNewThreadInput] = useState("");
+  const [files, setFiles] = useState<File[]>();
   const newThreadInputRef = useRef<ReactCodeMirrorRef>(null);
   const newMessageInputRef = useRef<ReactCodeMirrorRef>(null);
   const scrollContainerRef = useRef<HTMLDivElement>(null);
+  const fileInputRef = useRef<HTMLInputElement>(null);
   const messagesEndRef = useRef<HTMLDivElement>(null);
   const runtimeManager = useRuntimeManager();
   const { invokeAiTool } = useRequestClient();
 
+  const activeChatId = activeChat?.id;
+
   const {
     messages,
-    input,
-    setInput,
-    setMessages,
-    append,
-    handleSubmit,
+    sendMessage,
     error,
     status,
-    reload,
+    regenerate,
     stop,
+    addToolResult,
+    id: chatId,
   } = useChat({
-    id: activeChat?.id,
-    maxSteps: 10,
-    initialMessages: activeChat?.messages || [],
-    keepLastMessageOnError: true,
-    // Throttle the messages and data updates to 100ms
-    // experimental_throttle: 100,
-    api: runtimeManager.getAiURL("chat").toString(),
-    headers: runtimeManager.headers(),
-    experimental_prepareRequestBody: (options) => {
-      return {
-        ...options,
-        ...getAICompletionBody({
-          input: options.messages.map((m) => m.content).join("\n"),
-        }),
-        includeOtherCode: getCodes(""),
-      };
+    id: activeChatId,
+    // Only automatically submit if we have tool calls but no text response yet
+    sendAutomaticallyWhen: ({ messages }) => {
+      if (messages.length === 0) {
+        return false;
+      }
+
+      const lastMessage = messages[messages.length - 1];
+      const parts = lastMessage.parts;
+
+      if (parts.length === 0) {
+        return false;
+      }
+
+      // Only auto-send if the last message is an assistant message
+      // Because assistant messages are the ones that can have tool calls
+      if (lastMessage.role !== "assistant") {
+        return false;
+      }
+
+      const toolParts = parts.filter((part) =>
+        part.type.startsWith("tool-"),
+      ) as ToolUIPart[];
+
+      const hasCompletedToolCalls = toolParts.some(
+        (part) => part.state === "output-available",
+      );
+
+      // Check if the last part has any text content
+      const lastPart = parts[parts.length - 1];
+      const hasTextContent =
+        lastPart.type === "text" && lastPart.text?.trim().length > 0;
+
+      // Only auto-send if we have completed tool calls and there is no reply yet
+      return hasCompletedToolCalls && !hasTextContent;
     },
-    onFinish: (message) => {
+    messages: activeChat?.messages || [], // initial messages
+    transport: new DefaultChatTransport({
+      api: runtimeManager.getAiURL("chat").toString(),
+      headers: runtimeManager.headers(),
+      prepareSendMessagesRequest: async (options) => {
+        // Map from parts to a single string
+        function toContent(parts: UIMessage["parts"]): string {
+          return parts
+            .map((part) => (part.type === "text" ? part.text : ""))
+            .join("\n");
+        }
+
+        const input = toContent(options.messages.flatMap((m) => m.parts));
+        const completionBody = await getAICompletionBodyWithAttachments({
+          input,
+        });
+
+        // Map from UIMessage to our ChatMessage type
+        // If it's the last message, add the attachments from the completion body
+        function mapMessage(m: UIMessage, isLastMessage: boolean): ChatMessage {
+          const parts = m.parts;
+          if (isLastMessage) {
+            parts.push(...completionBody.attachments);
+          }
+          return {
+            role: m.role,
+            content: toContent(m.parts),
+            parts: parts,
+          };
+        }
+
+        return {
+          body: {
+            ...options,
+            ...completionBody.body,
+            messages: options.messages.map((m, idx) => ({
+              ...m,
+              ...mapMessage(m, idx === options.messages.length - 1),
+            })),
+          },
+        };
+      },
+    }),
+    onFinish: ({ messages }) => {
       setChatState((prev) => {
-        return addMessageToChat(
-          prev,
-          prev.activeChatId,
-          message.id,
-          "assistant",
-          message.content,
-          message.parts,
-        );
+        return replaceMessagesInChat({
+          chatState: prev,
+          chatId: prev.activeChatId,
+          messages: messages,
+        });
       });
     },
     onToolCall: async ({ toolCall }) => {
       try {
         const response = await invokeAiTool({
           toolName: toolCall.toolName,
-          arguments: toolCall.args as Record<string, never>,
+          arguments: toolCall.input as Record<string, never>,
+        });
+        addToolResult({
+          tool: toolCall.toolName,
+          toolCallId: toolCall.toolCallId,
+          output: response.result || response.error,
         });
-
-        // This response triggers the onFinish callback
-        return response.result || response.error;
       } catch (error) {
         Logger.error("Tool call failed:", error);
-        return `Error: ${error instanceof Error ? error.message : String(error)}`;
+        addToolResult({
+          tool: toolCall.toolName,
+          toolCallId: toolCall.toolCallId,
+          output: `Error: ${error instanceof Error ? error.message : String(error)}`,
+        });
       }
     },
     onError: (error) => {
       Logger.error("An error occurred:", error);
     },
-    onResponse: (response) => {
-      Logger.debug("Received HTTP response from server:", response);
-    },
   });
 
-  const isLoading = status === "submitted" || status === "streaming";
-
-  // Sync user messages from useChat to storage when they become available
-  // Only when we are done loading, for performance.
-  useEffect(() => {
-    if (!chatState.activeChatId || messages.length === 0 || isLoading) {
+  const onAddFiles = useEvent((files: File[]) => {
+    if (files.length === 0) {
       return;
     }
 
-    // Only sync if the last message is from a user
-    const lastMessage = messages[messages.length - 1];
-    if (lastMessage?.role !== "user") {
-      return;
+    let fileSize = 0;
+    for (const file of files) {
+      fileSize += file.size;
     }
 
-    const currentChat = chatState.chats.get(chatState.activeChatId);
-    if (!currentChat) {
+    if (fileSize > MAX_ATTACHMENT_SIZE) {
+      toast({
+        title: "File size exceeds 50MB limit",
+        description: "Please remove some files and try again.",
+      });
       return;
     }
 
-    const storedMessageIds = new Set(currentChat.messages.map((m) => m.id));
-
-    // Find user messages from useChat that aren't in storage yet
-    const missingUserMessages = messages.filter(
-      (m) => m.role === "user" && !storedMessageIds.has(m.id),
-    );
-
-    if (missingUserMessages.length > 0) {
-      setChatState((prev) => {
-        let result = prev;
-
-        for (const userMessage of missingUserMessages) {
-          result = addMessageToChat(
-            result,
-            prev.activeChatId,
-            userMessage.id,
-            "user",
-            userMessage.content,
-          );
-        }
+    setFiles((prev) => [...(prev ?? []), ...files]);
+  });
 
-        return result;
-      });
+  const removeFile = useEvent((file: File) => {
+    if (files) {
+      setFiles(files.filter((f) => f !== file));
     }
-  }, [
-    messages,
-    chatState.activeChatId,
-    chatState.chats,
-    setChatState,
-    isLoading,
-  ]);
+  });
+
+  const isLoading = status === "submitted" || status === "streaming";
 
   // Check if we're currently streaming reasoning in the latest message
   const isStreamingReasoning =
@@ -531,19 +720,19 @@ const ChatPanelBody = () => {
     };
 
     requestAnimationFrame(scrollToBottom);
-  }, [chatState.activeChatId]);
+  }, [activeChatId]);
 
-  const createNewThread = (initialMessage: string) => {
-    const CURRENT_TIME = Date.now();
+  const createNewThread = async (
+    initialMessage: string,
+    initialAttachments?: File[],
+  ) => {
+    const now = Date.now();
     const newChat: Chat = {
-      id: generateUUID() as ChatId,
-      title:
-        initialMessage.length > 50
-          ? `${initialMessage.slice(0, 50)}...`
-          : initialMessage,
-      messages: [], // Don't pre-populate - let useChat handle it and sync back
-      createdAt: CURRENT_TIME,
-      updatedAt: CURRENT_TIME,
+      id: chatId as ChatId,
+      title: generateChatTitle(initialMessage),
+      messages: [],
+      createdAt: now,
+      updatedAt: now,
     };
 
     // Create new chat and set as active
@@ -558,12 +747,23 @@ const ChatPanelBody = () => {
       return newState;
     });
 
+    const fileParts =
+      initialAttachments && initialAttachments.length > 0
+        ? await convertToFileUIPart(initialAttachments)
+        : undefined;
+
     // Trigger AI conversation with append
-    append({
-      id: generateUUID(),
+    sendMessage({
       role: "user",
-      content: initialMessage,
+      parts: [
+        {
+          type: "text" as const,
+          text: initialMessage,
+        },
+        ...(fileParts ?? []),
+      ],
     });
+    setFiles(undefined);
     setInput("");
   };
 
@@ -571,51 +771,43 @@ const ChatPanelBody = () => {
     setActiveChat(null);
     setInput("");
     setNewThreadInput("");
+    setFiles(undefined);
   });
 
   const handleMessageEdit = useEvent((index: number, newValue: string) => {
-    // Truncate both useChat and storage
-    setMessages((messages) => messages.slice(0, index));
-    const activeChatId = chatState.activeChatId;
-    if (activeChatId) {
-      setChatState((prev) => {
-        const nextChats = new Map(prev.chats);
-        const activeChat = chatState.chats.get(activeChatId);
-        if (activeChat) {
-          nextChats.set(activeChat.id, {
-            ...activeChat,
-            messages: activeChat.messages.slice(0, index),
-            updatedAt: Date.now(),
-          });
-        }
+    const editedMessage = messages[index];
+    const fileParts = editedMessage.parts?.filter((p) => p.type === "file");
 
-        return {
-          ...prev,
-          chats: nextChats,
-        };
-      });
-    }
-
-    append({
+    const messageId = editedMessage.id;
+    sendMessage({
+      messageId: messageId, // replace the message
       role: "user",
-      content: newValue,
+      parts: [{ type: "text", text: newValue }, ...fileParts],
     });
   });
 
   const handleChatInputSubmit = useEvent(
-    (e: KeyboardEvent | undefined, newValue: string): void => {
+    async (e: KeyboardEvent | undefined, newValue: string): Promise<void> => {
       if (!newValue.trim()) {
         return;
       }
       if (newMessageInputRef.current?.view) {
         storePrompt(newMessageInputRef.current.view);
       }
-      handleSubmit(e);
+      const fileParts = files ? await convertToFileUIPart(files) : undefined;
+
+      e?.preventDefault();
+      sendMessage({
+        text: newValue,
+        files: fileParts,
+      });
+      setInput("");
+      setFiles(undefined);
     },
   );
 
   const handleReload = () => {
-    reload();
+    regenerate();
   };
 
   const handleNewThreadSubmit = useEvent(() => {
@@ -625,16 +817,57 @@ const ChatPanelBody = () => {
     if (newThreadInputRef.current?.view) {
       storePrompt(newThreadInputRef.current.view);
     }
-    createNewThread(newThreadInput.trim());
+    createNewThread(newThreadInput.trim(), files);
   });
 
   const handleOnCloseThread = () => newThreadInputRef.current?.editor?.blur();
 
-  const sortedChats = useMemo(() => {
-    return [...chatState.chats.values()].sort(
-      (a, b) => b.updatedAt - a.updatedAt,
-    );
-  }, [chatState.chats]);
+  const isNewThread = messages.length === 0;
+  const chatInput = isNewThread ? (
+    <ChatInput
+      key="new-thread-input"
+      placeholder={`Ask anything, ${CONTEXT_TRIGGER} to include context about tables or dataframes`}
+      input={newThreadInput}
+      inputRef={newThreadInputRef}
+      inputClassName="px-1 py-0"
+      setInput={setNewThreadInput}
+      onSubmit={handleNewThreadSubmit}
+      isLoading={isLoading}
+      onStop={stop}
+      fileInputRef={fileInputRef}
+      onAddFiles={onAddFiles}
+      onClose={handleOnCloseThread}
+    />
+  ) : (
+    <ChatInput
+      input={input}
+      setInput={setInput}
+      onSubmit={handleChatInputSubmit}
+      inputRef={newMessageInputRef}
+      isLoading={isLoading}
+      onStop={stop}
+      onClose={() => newMessageInputRef.current?.editor?.blur()}
+      fileInputRef={fileInputRef}
+      onAddFiles={onAddFiles}
+    />
+  );
+
+  const filesPills = files && files.length > 0 && (
+    <div
+      className={cn(
+        "flex flex-row gap-1 flex-wrap p-1",
+        isNewThread && "py-2 px-1",
+      )}
+    >
+      {files?.map((file) => (
+        <FileAttachmentPill
+          file={file}
+          key={file.name}
+          onRemove={() => removeFile(file)}
+        />
+      ))}
+    </div>
+  );
 
   return (
     <div className="flex flex-col h-[calc(100%-53px)]">
@@ -643,7 +876,6 @@ const ChatPanelBody = () => {
           onNewChat={handleNewChat}
           activeChatId={activeChat?.id}
           setActiveChat={setActiveChat}
-          chats={sortedChats}
         />
       </TooltipProvider>
 
@@ -651,36 +883,19 @@ const ChatPanelBody = () => {
         className="flex-1 px-3 bg-(--slate-1) gap-4 py-3 flex flex-col overflow-y-auto"
         ref={scrollContainerRef}
       >
-        {(!messages || messages.length === 0) && (
-          <div className="flex flex-col rounded-md border bg-background">
-            <div className="px-1">
-              <PromptInput
-                inputRef={newThreadInputRef}
-                key="new-thread-input"
-                value={newThreadInput}
-                placeholder="Ask anything, @ to include context about tables or dataframes"
-                onClose={handleOnCloseThread}
-                onChange={setNewThreadInput}
-                onSubmit={handleNewThreadSubmit}
-              />
-            </div>
-            <ChatInputFooter
-              isEmpty={!newThreadInput.trim()}
-              onSendClick={handleNewThreadSubmit}
-              isLoading={isLoading}
-              onStop={stop}
-            />
+        {isNewThread && (
+          <div className="rounded-md border bg-background">
+            {filesPills}
+            {chatInput}
           </div>
         )}
 
         {messages.map((message, idx) => (
-          <ChatMessage
+          <ChatMessageDisplay
             key={message.id}
             message={message}
             index={idx}
             onEdit={handleMessageEdit}
-            setChatState={setChatState}
-            chatState={chatState}
             isStreamingReasoning={isStreamingReasoning}
             isLast={idx === messages.length - 1}
           />
@@ -712,40 +927,15 @@ const ChatPanelBody = () => {
         </div>
       )}
 
-      {messages && messages.length > 0 && (
-        <ChatInput
-          input={input}
-          setInput={setInput}
-          onSubmit={handleChatInputSubmit}
-          inputRef={newMessageInputRef}
-          isLoading={isLoading}
-          onStop={stop}
-        />
+      {/* For existing threads, we place the chat input at the bottom */}
+      {!isNewThread && (
+        <>
+          {filesPills}
+          {chatInput}
+        </>
       )}
     </div>
   );
 };
 
-function isLastMessageReasoning(messages: Message[]): boolean {
-  if (messages.length === 0) {
-    return false;
-  }
-
-  const lastMessage = messages.at(-1);
-  if (!lastMessage) {
-    return false;
-  }
-
-  if (lastMessage.role !== "assistant" || !lastMessage.parts) {
-    return false;
-  }
-
-  const parts = lastMessage.parts;
-  if (parts.length === 0) {
-    return false;
-  }
-
-  // Check if the last part is reasoning
-  const lastPart = parts[parts.length - 1];
-  return lastPart.type === "reasoning";
-}
+export default ChatPanel;