@marimo-team/islands 0.25.3-dev34 → 0.25.3-dev348

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

@@ -13,28 +13,13 @@ import {
 import { useAtom, useAtomValue, useSetAtom, useStore } from "jotai";
 import {
   BotMessageSquareIcon,
-  HatGlasses,
-  ArrowRightIcon,
   Loader2,
-  type LucideIcon,
-  MessageCircleIcon,
-  NotebookText,
   PlusIcon,
-  SparklesIcon,
   SettingsIcon,
-  CodeIcon,
 } from "lucide-react";
-import { memo, useEffect, 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 {
-  Select,
-  SelectContent,
-  SelectGroup,
-  SelectItem,
-  SelectLabel,
-  SelectTrigger,
-} from "@/components/ui/select";
 import { replaceMessagesInChat } from "@/core/ai/chat-utils";
 import { useAIConfigActions } from "@/core/ai/config";
 import { AI_SDK_UI_THROTTLE_MS } from "@/core/ai/constants";
@@ -49,37 +34,33 @@ import {
   pendingAiPromptAtom,
 } from "@/core/ai/state";
 import type { ToolNotebookContext } from "@/core/ai/tools/base";
-import {
-  type CopilotMode,
-  FRONTEND_TOOL_REGISTRY,
-} from "@/core/ai/tools/registry";
-import { useCellActions } from "@/core/cells/cells";
+import { FRONTEND_TOOL_REGISTRY } from "@/core/ai/tools/registry";
+import { useCellActions } from "@/core/cells/actions";
 import { aiAtom, aiModelConfiguredAtom } from "@/core/config/config";
 import { DEFAULT_AI_MODEL } from "@/core/config/config-schema";
 import { useRequestClient } from "@/core/network/requests";
 import { useRuntimeManager } from "@/core/runtime/config";
 import { isWasm } from "@/core/wasm/utils";
 import { ErrorBanner } from "@/plugins/impl/common/error-banner";
-import { cn } from "@/utils/cn";
 import { Logger } from "@/utils/Logger";
 import { AIModelDropdown } from "../ai/ai-model-dropdown";
 import { useOpenSettingsToTab } from "../app-config/state";
 import {
   PairWithAgentBanner,
-  PairWithAgentButton,
+  usePairWithAgentModal,
 } from "../editor/actions/pair-with-agent";
-import { PairWithAgentModal } from "../editor/actions/pair-with-agent-modal";
-import { PromptInput } from "../editor/ai/add-cell-with-ai";
+import { PromptInput } from "../editor/ai/prompt-input";
+import { PendingAICells } from "../editor/chrome/wrapper/pending-ai-cells";
 import {
   addContextCompletion,
-  CONTEXT_TRIGGER,
   isContextAttachment,
   resolveChatContext,
 } from "../editor/ai/completion-utils";
 import { StreamingChunkTransport } from "../editor/ai/transport/chat-transport";
+import { PanelHeaderActions } from "../editor/chrome/panels/panel-header";
 import { PanelEmptyState } from "../editor/chrome/panels/empty-state";
 import { CopyClipboardIcon } from "../icons/copy-icon";
-import { useImperativeModal } from "../modal/ImperativeModal";
+import "./chat-typography.css";
 import { MCPStatusIndicator } from "../mcp/mcp-status-indicator";
 import { Tooltip, TooltipProvider } from "../ui/tooltip";
 import {
@@ -94,9 +75,11 @@ import {
   getAbortReasonFromChunk,
   resolveChatAbortReason,
 } from "./chat-abort";
-import { renderUIMessage } from "./chat-display";
+import { AgentActivity } from "./agent-activity";
+import { getAnswerText, renderUIMessage } from "./chat-display";
 import { ChatHistoryPopover } from "./chat-history-popover";
 import { CapabilitiesPopover } from "./capabilities-popover";
+import { NotebookAccessSelector } from "./notebook-access-selector";
 import {
   type ChatMessagePart,
   convertToFileUIPart,
@@ -104,20 +87,24 @@ import {
   handleToolCall,
   hasPendingToolCalls,
   hasUnresolvedToolCalls,
-  isLastMessageReasoning,
   PROVIDERS_THAT_SUPPORT_ATTACHMENTS,
   type QueuedUserMessage,
   shouldFlushQueue,
   useFileState,
   useMessageQueue,
 } from "./chat-utils";
-import { getCodes } from "@/core/codemirror/copilot/getCodes";
+import { notebookCellCodes } from "@/core/codemirror/copilot/getCodes";
+import { getSessionId } from "@/core/kernel/session";
+import { withNotebookContext } from "./notebook-context";
 import { focusInputAndMoveToEnd } from "@/core/codemirror/utils";
 import ScrollToBottomButton from "./acp/scroll-to-bottom-button";
 import { useChatControllerId } from "./use-chat-controller-id";
+import { isChatBusy, useChatGeneration } from "./use-chat-generation";
+import { useAssistantTurn } from "./use-assistant-turn";
+import { UserMessage } from "./user-message";
 
 // Default mode for the AI
-const DEFAULT_MODE = "manual";
+const DEFAULT_MODE = "ask";
 
 const QueuedMessageDisplay = memo(
   ({ message }: { message: QueuedUserMessage }) => {
@@ -131,7 +118,7 @@ const QueuedMessageDisplay = memo(
 
     return (
       <div className="flex justify-end">
-        <div className="w-[95%] bg-background border border-dashed p-2 rounded-sm opacity-70">
+        <div className="w-[95%] bg-background border border-dashed p-2 rounded-[10px] opacity-70">
           {fileParts.map((filePart, idx) => (
             <AttachmentRenderer attachment={filePart} key={idx} />
           ))}
@@ -164,35 +151,41 @@ const ChatHeader: React.FC<ChatHeaderProps> = ({
   const { handleClick } = useOpenSettingsToTab();
 
   return (
-    <div className="flex border-b px-2 py-1 justify-between shrink-0 items-center">
-      <Tooltip content="New chat">
-        <Button variant="text" size="icon" onClick={onNewChat}>
-          <PlusIcon className="h-4 w-4" />
-        </Button>
-      </Tooltip>
-      <div className="flex items-center gap-2">
-        <Tooltip content="Connect your agent">
-          <PairWithAgentButton className="mr-1" label="Connect" />
-        </Tooltip>
-        <MCPStatusIndicator />
-        <Tooltip content="AI Settings">
+    <PanelHeaderActions>
+      <div className="mo-ai-panel flex w-full shrink-0 items-center font-sans text-sm">
+        <div className="flex w-full items-center gap-1">
           <Button
-            variant="text"
-            size="xs"
-            className="hover:bg-foreground/10 py-2"
-            onClick={() => handleClick("ai")}
+            variant="ghost"
+            size="sm"
+            className="chat-control gap-1.5 font-normal"
+            onClick={onNewChat}
           >
-            <SettingsIcon className="h-4 w-4" />
+            <PlusIcon className="size-4" aria-hidden={true} />
+            New chat
           </Button>
-        </Tooltip>
-        <ChatHistoryPopover
-          activeChatId={activeChatId}
-          setActiveChat={setActiveChat}
-          onDeleteChats={onDeleteChats}
-          isLoading={isLoading}
-        />
+          <ChatHistoryPopover
+            activeChatId={activeChatId}
+            setActiveChat={setActiveChat}
+            onDeleteChats={onDeleteChats}
+            isLoading={isLoading}
+          />
+          <div className="ml-auto flex items-center gap-2">
+            <MCPStatusIndicator />
+            <Tooltip content="AI Settings">
+              <Button
+                className="chat-control"
+                aria-label="AI Settings"
+                variant="ghost"
+                size="icon"
+                onClick={() => handleClick("ai")}
+              >
+                <SettingsIcon className="size-4" aria-hidden={true} />
+              </Button>
+            </Tooltip>
+          </div>
+        </div>
       </div>
-    </div>
+    </PanelHeaderActions>
   );
 };
 
@@ -200,9 +193,9 @@ interface ChatMessageProps {
   message: UIMessage;
   index: number;
   onEdit: (index: number, newValue: string) => void;
-  isStreamingReasoning: boolean;
   isLast: boolean;
   isActive: boolean;
+  turnStartedAt: number | null;
   addToolApprovalResponse?: ChatAddToolApproveResponseFunction;
 }
 
@@ -211,78 +204,40 @@ const ChatMessageDisplay: React.FC<ChatMessageProps> = memo(
     message,
     index,
     onEdit,
-    isStreamingReasoning,
     isLast,
     isActive,
+    turnStartedAt,
     addToolApprovalResponse,
   }) => {
-    const renderUserMessage = (message: UIMessage) => {
-      const textParts = message.parts?.filter(
-        (p): p is TextUIPart => p.type === "text",
-      );
-      const content = textParts?.map((p) => p.text).join("\n");
-      const fileParts = message.parts?.filter(
-        (p): p is FileUIPart => p.type === "file",
-      );
-
+    if (message.role === "user") {
       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={content}
-            placeholder="Type your message..."
-            onChange={() => {
-              // noop
-            }}
-            onSubmit={(_e, newValue) => {
-              if (!newValue.trim()) {
-                return;
-              }
-              onEdit(index, newValue);
-            }}
-            onClose={() => {
-              // noop
-            }}
+        <div className="group flex justify-end">
+          <UserMessage
+            message={message}
+            canEdit={!isActive}
+            onEdit={(newValue) => onEdit(index, newValue)}
           />
         </div>
       );
-    };
-
-    const renderOtherMessage = (message: UIMessage) => {
-      const textParts = message.parts.filter(
-        (p): p is TextUIPart => p.type === "text",
-      );
-      const content = textParts.map((p) => p.text).join("\n");
+    }
 
-      return (
+    return (
+      <div className="group relative flex justify-start">
         <div className="w-[95%] wrap-break-word">
-          <div className="absolute right-1 top-1 opacity-0 group-hover:opacity-100 transition-opacity">
-            <CopyClipboardIcon className="h-3 w-3" value={content || ""} />
+          <div className="absolute right-1 top-1 opacity-0 group-hover:opacity-100 group-focus-within:opacity-100 transition-opacity">
+            <CopyClipboardIcon
+              className="h-3 w-3"
+              value={() => getAnswerText(message)}
+            />
           </div>
           {renderUIMessage({
             message,
-            isStreamingReasoning,
             isLast,
             isActive,
+            turnStartedAt,
             addToolApprovalResponse,
           })}
         </div>
-      );
-    };
-
-    return (
-      <div
-        className={cn(
-          "flex group relative",
-          message.role === "user" ? "justify-end" : "justify-start",
-        )}
-      >
-        {message.role === "user"
-          ? renderUserMessage(message)
-          : renderOtherMessage(message)}
       </div>
     );
   },
@@ -316,92 +271,27 @@ const ChatInputFooter: React.FC<ChatInputFooterProps> = memo(
 
     const { saveModeChange } = useAIConfigActions();
 
-    const modeOptions: {
-      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: "AI with access to read-only tools like documentation search",
-        Icon: NotebookText,
-      },
-      {
-        value: "agent",
-        label: "Agent",
-        subtitle: "AI with access to read and write tools",
-        Icon: HatGlasses,
-      },
-      {
-        value: "code_mode",
-        label: "Code Mode (beta)",
-        subtitle:
-          "AI with access to the notebook's kernel. Potentially destructive.",
-        Icon: CodeIcon,
-      },
-    ];
-
     const isAttachmentSupported =
       PROVIDERS_THAT_SUPPORT_ATTACHMENTS.has(currentProvider);
 
-    const currentModeOption = modeOptions.find((o) => o.value === currentMode);
-    const CurrentModeIcon = currentModeOption?.Icon;
-    const CurrentModeLabel = currentModeOption?.label;
-
     return (
       <TooltipProvider>
-        <div className="px-3 py-2 border-t border-border/20 flex flex-wrap items-center gap-1">
-          <div className="flex flex-wrap items-center gap-1">
-            <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 shrink-0">
-                {CurrentModeIcon && <CurrentModeIcon className="h-3 w-3" />}
-                <span className="text-nowrap">{CurrentModeLabel}</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}
-                          </span>
-                        </div>
-                      </div>
-                    </SelectItem>
-                  ))}
-                </SelectGroup>
-              </SelectContent>
-            </Select>
+        <div className="chat-input-footer px-3 py-2 border-t border-border/20 flex items-center gap-2">
+          <div className="chat-input-selectors flex min-w-0 flex-1 items-center gap-1">
+            <NotebookAccessSelector
+              value={currentMode}
+              onValueChange={saveModeChange}
+            />
             <AIModelDropdown
               placeholder="Model"
-              triggerClassName="h-6 text-xs shadow-none! ring-0! bg-muted hover:bg-muted/30 rounded-sm max-w-[200px]"
+              triggerClassName="chat-control chat-selector h-6 text-xs rounded-sm min-w-0 max-w-[200px] flex-1"
               iconSize="small"
               showAddCustomModelDocs={true}
               forRole="chat"
             />
-            <CapabilitiesPopover />
           </div>
-          <div className="flex flex-row ml-auto">
+          <div className="chat-input-actions flex shrink-0 items-center gap-0.5 ml-auto">
+            <CapabilitiesPopover />
             <AddContextButton
               handleAddContext={onAddContext}
               isLoading={isLoading}
@@ -429,9 +319,7 @@ 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>;
@@ -444,9 +332,7 @@ interface ChatInputProps {
 
 const ChatInput: React.FC<ChatInputProps> = memo(
   ({
-    placeholder,
     input,
-    inputClassName,
     setInput,
     onSubmit,
     inputRef,
@@ -456,6 +342,20 @@ const ChatInput: React.FC<ChatInputProps> = memo(
     onAddFiles,
     onClose,
   }) => {
+    const openPairWithAgentModal = usePairWithAgentModal();
+    const placeholderAction = useMemo(
+      () =>
+        isWasm()
+          ? undefined
+          : {
+              beforeText: "Type your message or ",
+              linkText: "pair your own coding agent",
+              afterText: ". Use @ to reference variables and cells.",
+              onClick: openPairWithAgentModal,
+            },
+      [openPairWithAgentModal],
+    );
+
     const handleSendClick = useEvent(() => {
       if (input.trim()) {
         onSubmit(undefined, input);
@@ -463,8 +363,8 @@ const ChatInput: React.FC<ChatInputProps> = memo(
     });
 
     return (
-      <div className="relative shrink-0 flex flex-col border-t">
-        <div className={cn("px-2 py-1.5", inputClassName)}>
+      <div className="relative shrink-0 flex flex-col">
+        <div className="px-2 py-1.5">
           <PromptInput
             className="max-h-[400px]"
             inputRef={inputRef}
@@ -473,7 +373,8 @@ const ChatInput: React.FC<ChatInputProps> = memo(
             onSubmit={onSubmit}
             onClose={onClose}
             onAddFiles={onAddFiles}
-            placeholder={placeholder || "Type your message..."}
+            placeholder="Type your message. Use @ to reference variables and cells."
+            placeholderAction={placeholderAction}
           />
         </div>
         <ChatInputFooter
@@ -492,26 +393,6 @@ const ChatInput: React.FC<ChatInputProps> = memo(
 
 ChatInput.displayName = "ChatInput";
 
-const PairWithAgentCallout: React.FC<{
-  onPairWithAgent: () => void;
-}> = ({ onPairWithAgent }) => {
-  if (isWasm()) {
-    return null;
-  }
-
-  return (
-    <Button
-      variant="text"
-      className="gap-1.5 text-sm text-link hover:underline"
-      onClick={onPairWithAgent}
-    >
-      <SparklesIcon className="h-3.5 w-3.5 shrink-0" />
-      <span>Connect your own agent to this notebook</span>
-      <ArrowRightIcon className="h-3 w-3 shrink-0" />
-    </Button>
-  );
-};
-
 const ChatPanel = () => {
   const aiConfigured = useAtomValue(aiModelConfiguredAtom);
   const { handleClick } = useOpenSettingsToTab();
@@ -567,9 +448,10 @@ const ChatPanelBody = () => {
     lastAbortReasonRef.current = null;
     setAbortNotice(null);
   });
+  const turn = useAssistantTurn();
+  const resetTurn = turn.reset;
   const runtimeManager = useRuntimeManager();
-  const { invokeAiTool, sendRun } = useRequestClient();
-  const { openModal, closeModal } = useImperativeModal();
+  const { invokeAiTool } = useRequestClient();
 
   const activeChatId = activeChat?.id;
   const { chatControllerId, renewDraftChatId } =
@@ -577,12 +459,10 @@ const ChatPanelBody = () => {
   const store = useStore();
 
   const { addStagedCell } = useStagedAICellsActions();
-  const { createNewCell, prepareForRun } = useCellActions();
+  const { createCell } = useCellActions();
   const toolContext: ToolNotebookContext = {
     addStagedCell,
-    createNewCell,
-    prepareForRun,
-    sendRun,
+    createCell,
     store,
   };
 
@@ -592,6 +472,7 @@ const ChatPanelBody = () => {
     error,
     status,
     regenerate,
+    setMessages,
     stop,
     addToolOutput,
     addToolApprovalResponse,
@@ -599,7 +480,8 @@ const ChatPanelBody = () => {
   } = useChat({
     id: chatControllerId,
     throttle: AI_SDK_UI_THROTTLE_MS,
-    sendAutomaticallyWhen: ({ messages }) => hasPendingToolCalls(messages),
+    sendAutomaticallyWhen: ({ messages }) =>
+      !turn.hasEnded && hasPendingToolCalls(messages),
     messages: activeChat?.messages || [], // initial messages
     transport: new StreamingChunkTransport(
       {
@@ -621,7 +503,6 @@ const ChatPanelBody = () => {
 
           const completionBody = {
             uiMessages: options.messages,
-            includeOtherCode: getCodes(""),
             options: store.get(chatOptionsAtom),
           };
 
@@ -646,7 +527,26 @@ const ChatPanelBody = () => {
         }
       },
     ),
-    onFinish: ({ messages, isError, isAbort, finishReason }) => {
+    onFinish: ({
+      messages: finishedMessages,
+      isError,
+      isAbort,
+      finishReason,
+    }) => {
+      // useChat routes every chat's callbacks to the latest one. Its
+      // setMessages isn't a React setter: the updater runs synchronously.
+      const finished: { messages?: typeof finishedMessages } = {};
+      setMessages((current) => {
+        if (current.at(-1)?.id !== finishedMessages.at(-1)?.id) {
+          return current;
+        }
+        finished.messages = turn.finish(finishedMessages, { isAbort, isError });
+        return finished.messages;
+      });
+      const { messages } = finished;
+      if (!messages) {
+        return;
+      }
       setChatState((prev) => {
         return replaceMessagesInChat({
           chatState: prev,
@@ -689,9 +589,21 @@ const ChatPanelBody = () => {
     },
   });
 
+  const contextualize = useEvent(
+    (parts: ChatMessagePart[], history: UIMessage[]) =>
+      withNotebookContext(parts, history, {
+        sessionId: getSessionId(),
+        cells:
+          store.get(aiAtom)?.mode === "code_mode"
+            ? undefined
+            : store.get(notebookCellCodes),
+      }),
+  );
+
   const sendUserMessage = useEvent((parts: ChatMessagePart[]) => {
     clearAbortState();
-    sendMessage({ role: "user", parts });
+    turn.start();
+    sendMessage({ role: "user", parts: contextualize(parts, messages) });
   });
 
   const tryFlushQueuedMessages = useEvent(
@@ -715,15 +627,43 @@ const ChatPanelBody = () => {
     },
   );
 
-  const isLoading = status === "submitted" || status === "streaming";
-  // Read via a ref so the queue-vs-send decision stays correct even when it is
-  // made after an `await` (e.g. resolving @-context), by which point the render
-  // closure's `isLoading` may be stale.
-  const isLoadingRef = useRef(isLoading);
-  isLoadingRef.current = isLoading;
+  const isLoading = isChatBusy(status);
+  const isTurnActive = turn.isActive({
+    isLoading,
+    hasError: error != null,
+    messages,
+  });
+  useChatGeneration(isTurnActive);
+
+  const stopTurn = useEvent(() => {
+    if (isLoading) {
+      turn.stop();
+      void stop();
+      return;
+    }
+    // Between requests there is no stream to abort, so `onFinish` won't run.
+    const stopped = turn.finish(messages, { isAbort: true, isError: false });
+    setMessages(stopped);
+    setChatState((prev) =>
+      replaceMessagesInChat({
+        chatState: prev,
+        chatId: prev.activeChatId,
+        messages: stopped,
+      }),
+    );
+    tryFlushQueuedMessages(stopped, { isError: false, isAbort: true });
+  });
 
+  const respondToApproval = useEvent(
+    (response: Parameters<ChatAddToolApproveResponseFunction>[0]) => {
+      turn.resume(messages.at(-1));
+      return addToolApprovalResponse(response);
+    },
+  );
+  // useEvent reads the latest loading state, even when called after an await
+  // (e.g. resolving @-context) that outlives the original render's closure.
   const submitOrQueue = useEvent((parts: ChatMessagePart[]) => {
-    if (isLoadingRef.current || hasQueuedRef.current) {
+    if (isTurnActive || hasQueuedRef.current) {
       enqueueUserMessage(parts);
     } else {
       sendUserMessage(parts);
@@ -756,10 +696,6 @@ const ChatPanelBody = () => {
     });
   });
 
-  // Check if we're currently streaming reasoning in the latest message
-  const isStreamingReasoning =
-    isLoading && messages.length > 0 && isLastMessageReasoning(messages);
-
   // Pin to the bottom while the user is already there.
   useEffect(() => {
     setIsScrolledToBottom(true);
@@ -767,6 +703,11 @@ const ChatPanelBody = () => {
     clearAbortState();
   }, [activeChatId, clearQueuedMessages, clearAbortState]);
 
+  // Creating a chat keeps the draft's controller, so its turn survives.
+  useEffect(() => {
+    resetTurn();
+  }, [chatControllerId, resetTurn]);
+
   useEffect(() => {
     if (!isScrolledToBottom) {
       return;
@@ -821,14 +762,18 @@ const ChatPanelBody = () => {
       const { contextPart, attachments } =
         await resolveChatContext(initialMessage);
 
+      turn.start();
       sendMessage({
         role: "user",
-        parts: [
-          { type: "text" as const, text: initialMessage },
-          ...(contextPart ? [contextPart] : []),
-          ...(fileParts ?? []),
-          ...attachments,
-        ],
+        parts: contextualize(
+          [
+            { type: "text" as const, text: initialMessage },
+            ...(contextPart ? [contextPart] : []),
+            ...(fileParts ?? []),
+            ...attachments,
+          ],
+          [],
+        ),
       });
       clearFiles();
       setInput("");
@@ -865,10 +810,6 @@ const ChatPanelBody = () => {
     });
   });
 
-  const handlePairWithAgent = useEvent(() => {
-    openModal(<PairWithAgentModal onClose={closeModal} />);
-  });
-
   const handleMessageEdit = useEvent(
     async (index: number, newValue: string) => {
       const editedMessage = messages[index];
@@ -882,15 +823,19 @@ const ChatPanelBody = () => {
       const { contextPart, attachments } = await resolveChatContext(newValue);
 
       const messageId = editedMessage.id;
+      turn.start();
       sendMessage({
         messageId: messageId, // replace the message
         role: "user",
-        parts: [
-          { type: "text", text: newValue },
-          ...(contextPart ? [contextPart] : []),
-          ...userFileParts,
-          ...attachments,
-        ],
+        parts: contextualize(
+          [
+            { type: "text", text: newValue },
+            ...(contextPart ? [contextPart] : []),
+            ...userFileParts,
+            ...attachments,
+          ],
+          messages.slice(0, index),
+        ),
       });
     },
   );
@@ -920,6 +865,7 @@ const ChatPanelBody = () => {
 
   const handleReload = () => {
     clearAbortState();
+    turn.start();
     regenerate();
   };
 
@@ -959,14 +905,20 @@ const ChatPanelBody = () => {
       return;
     }
     setPendingPrompt(null);
-    const { prompt, submit } = pendingPrompt;
+    const { prompt, submit, append } = pendingPrompt;
+    const applyPrompt = (draft: string) =>
+      append
+        ? draft.includes(prompt.trim())
+          ? draft
+          : `${draft}${draft ? " " : ""}${prompt}`
+        : prompt;
     if (submit) {
       void submitPendingPrompt(prompt);
     } else if (isNewThread) {
-      setNewThreadInput(prompt);
+      setNewThreadInput(applyPrompt);
       focusInputAndMoveToEnd(newThreadInputRef);
     } else {
-      setInput(prompt);
+      setInput(applyPrompt);
       focusInputAndMoveToEnd(newMessageInputRef);
     }
   }, [pendingPrompt, setPendingPrompt, isNewThread, submitPendingPrompt]);
@@ -974,14 +926,12 @@ const ChatPanelBody = () => {
   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}
+      isLoading={isTurnActive}
+      onStop={stopTurn}
       fileInputRef={fileInputRef}
       onAddFiles={addFiles}
       onClose={handleOnCloseThread}
@@ -992,8 +942,8 @@ const ChatPanelBody = () => {
       setInput={setInput}
       onSubmit={handleChatInputSubmit}
       inputRef={newMessageInputRef}
-      isLoading={isLoading}
-      onStop={stop}
+      isLoading={isTurnActive}
+      onStop={stopTurn}
       onClose={() => newMessageInputRef.current?.editor?.blur()}
       fileInputRef={fileInputRef}
       onAddFiles={addFiles}
@@ -1001,12 +951,7 @@ const ChatPanelBody = () => {
   );
 
   const filesPills = files && files.length > 0 && (
-    <div
-      className={cn(
-        "flex flex-row gap-1 flex-wrap p-1",
-        isNewThread && "py-2 px-1",
-      )}
-    >
+    <div className="flex flex-row gap-1 flex-wrap px-3 pt-2">
       {files?.map((file) => (
         <FileAttachmentPill
           file={file}
@@ -1017,8 +962,15 @@ const ChatPanelBody = () => {
     </div>
   );
 
+  const composer = (
+    <div className="@container rounded-[10px] border border-border bg-background">
+      {filesPills}
+      {chatInput}
+    </div>
+  );
+
   return (
-    <div className="flex flex-col h-[calc(100%-53px)]">
+    <div className="mo-ai-panel flex flex-col flex-1 min-h-0 font-prose text-sm leading-6">
       <TooltipProvider>
         <ChatHeader
           onNewChat={handleNewChat}
@@ -1031,43 +983,44 @@ const ChatPanelBody = () => {
 
       <div className="flex-1 flex flex-col overflow-hidden relative min-h-0">
         <div
-          className="flex-1 px-3 bg-(--slate-1) gap-4 py-3 flex flex-col overflow-y-auto"
+          className="flex-1 px-3 bg-background gap-4 py-3 flex flex-col overflow-y-auto"
           ref={scrollContainerRef}
           onScroll={handleScroll}
         >
-          {isNewThread && (
-            <div className="flex flex-col gap-2">
-              <div className="rounded-md border bg-background">
-                {filesPills}
-                {chatInput}
-              </div>
-              <PairWithAgentCallout onPairWithAgent={handlePairWithAgent} />
-            </div>
-          )}
-
-          {messages.map((message, idx) => (
-            <ChatMessageDisplay
-              key={message.id}
-              message={message}
-              index={idx}
-              onEdit={handleMessageEdit}
-              isStreamingReasoning={isStreamingReasoning}
-              isLast={idx === messages.length - 1}
-              isActive={isLoading}
-              addToolApprovalResponse={addToolApprovalResponse}
+          <TooltipProvider>
+            {messages.map((message, idx) => {
+              const isLast = idx === messages.length - 1;
+              return (
+                <ChatMessageDisplay
+                  key={message.id}
+                  message={message}
+                  index={idx}
+                  onEdit={handleMessageEdit}
+                  isLast={isLast}
+                  isActive={isTurnActive}
+                  turnStartedAt={isLast ? turn.startedAt : null}
+                  addToolApprovalResponse={respondToApproval}
+                />
+              );
+            })}
+          </TooltipProvider>
+
+          {isTurnActive && messages.at(-1)?.role === "user" && (
+            <AgentActivity
+              running={true}
+              thinking={true}
+              timing={
+                turn.startedAt == null
+                  ? undefined
+                  : { startedAt: turn.startedAt }
+              }
             />
-          ))}
+          )}
 
           {queuedMessages.map((message) => (
             <QueuedMessageDisplay key={message.id} message={message} />
           ))}
 
-          {isLoading && (
-            <div className="flex justify-center py-4">
-              <Loader2 className="h-4 w-4 animate-spin" />
-            </div>
-          )}
-
           {error && (
             <div className="flex items-center justify-center space-x-2 mb-4">
               <ErrorBanner error={error || new Error("Unknown error")} />
@@ -1095,21 +1048,8 @@ const ChatPanelBody = () => {
         />
       </div>
 
-      {isLoading && (
-        <div className="w-full flex justify-center items-center z-20 border-t">
-          <Button variant="linkDestructive" size="sm" onClick={stop}>
-            Stop
-          </Button>
-        </div>
-      )}
-
-      {/* For existing threads, we place the chat input at the bottom */}
-      {!isNewThread && (
-        <>
-          {filesPills}
-          {chatInput}
-        </>
-      )}
+      {!isTurnActive && <PendingAICells />}
+      <div className="shrink-0 p-3">{composer}</div>
     </div>
   );
 };