@marimo-team/islands 0.25.3-dev266 → 0.25.3-dev268

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

@@ -7,103 +7,95 @@ import type {
   ToolUIPart,
   UIMessage,
 } from "ai";
-import { ExternalLinkIcon, FileTextIcon } from "lucide-react";
+import { BrainIcon, ExternalLinkIcon, FileTextIcon } from "lucide-react";
 import React from "react";
 import { renderHTML } from "@/plugins/core/RenderHTML";
 import { logNever } from "@/utils/assertNever";
-import { Logger } from "@/utils/Logger";
-import { ChatMarkdown } from "./chat-markdown";
+import { AgentActivity } from "./agent-activity";
 import { AttachmentRenderer, SourceChip } from "./chat-components";
-import { ReasoningAccordion } from "./reasoning-accordion";
-import { ToolActivityGroup } from "./tool-call/tool-activity-group";
+import { ChatMarkdown } from "./chat-markdown";
+import { PENDING_TOOL_STATES } from "./tool-call/shared";
 import { ToolCallView } from "./tool-call/tool-call-view";
+import { getTurnTiming } from "./turn-timing";
+
+type Part = UIMessage["parts"][number];
+interface IndexedPart {
+  part: Part;
+  index: number;
+}
 
 export const renderUIMessage = ({
   message,
-  isStreamingReasoning,
   isLast,
   isActive,
+  turnStartedAt,
   addToolApprovalResponse,
 }: {
   message: UIMessage;
-  isStreamingReasoning: boolean;
   isLast: boolean;
   /** Whether the chat is currently streaming/submitting a response. */
   isActive: boolean;
+  turnStartedAt?: number | null;
   addToolApprovalResponse?: ChatAddToolApproveResponseFunction;
-}) => {
+}): React.ReactNode[] => {
+  const running = isActive && isLast;
+  const { work, rest } = splitTurn(message.parts);
   const content: React.ReactNode[] = [];
-  let activity: { part: ToolUIPart | DynamicToolUIPart; index: number }[] = [];
 
-  const flushActivity = () => {
-    if (activity.length === 0) {
-      return;
-    }
+  if (work.length > 0 || (running && rest.length === 0)) {
+    const tools = work.flatMap(({ part }) =>
+      isAnyToolPart(part) ? [part] : [],
+    );
+    const timing =
+      running && turnStartedAt != null
+        ? { startedAt: turnStartedAt }
+        : getTurnTiming(message);
     content.push(
-      <ToolActivityGroup
-        key={activity[0].index}
-        tools={activity.map(({ part }) => ({
-          toolName: part.type === "dynamic-tool" ? part.toolName : part.type,
-          state: part.state,
-        }))}
-        isActive={isActive && isLast}
-      >
-        {activity.map(({ part, index }) => renderUIMessagePart(part, index))}
-      </ToolActivityGroup>,
+      <AgentActivity
+        key="activity"
+        running={running}
+        thinking={
+          rest.length === 0 &&
+          (work.length === 0 || work.at(-1)?.part.type === "reasoning")
+        }
+        interrupted={
+          !running &&
+          (timing?.stopped === true ||
+            tools.some((part) => PENDING_TOOL_STATES.has(part.state)))
+        }
+        defaultOpen={!rest.some(({ part }) => !isPendingApproval(part))}
+        timing={timing}
+        toolCount={tools.length}
+        failedCount={
+          tools.filter((part) => part.state === "output-error").length
+        }
+        steps={work.map(({ part, index }) => renderPart(part, index, true))}
+      />,
     );
-    activity = [];
-  };
-
-  message.parts.forEach((part, index) => {
-    // SDK step boundaries and invisible parts should not split visible activity.
-    if (
-      part.type === "step-start" ||
-      isDataPart(part) ||
-      (part.type === "text" && !part.text.trim())
-    ) {
-      return;
-    }
-    if (
-      (isToolPart(part) || part.type === "dynamic-tool") &&
-      part.state !== "output-error" &&
-      part.state !== "output-denied" &&
-      part.state !== "approval-requested"
-    ) {
-      activity.push({ part, index });
-      return;
-    }
-    flushActivity();
-    content.push(renderUIMessagePart(part, index));
-  });
-  flushActivity();
-
+  }
+  for (const { part, index } of rest) {
+    content.push(renderPart(part, index, false));
+  }
   return content;
 
-  function renderUIMessagePart(
-    part: UIMessage["parts"][number],
-    index: number,
-  ) {
-    if (isToolPart(part)) {
+  function renderPart(part: Part, index: number, inActivity: boolean) {
+    if (isToolPart(part) || part.type === "dynamic-tool") {
       return (
         <ToolCallView
           key={index}
-          index={index}
-          toolName={part.type}
+          toolName={part.type === "dynamic-tool" ? part.toolName : part.type}
           result={part.output}
           errorText={part.state === "output-error" ? part.errorText : undefined}
-          className="my-1"
           state={part.state}
           input={part.input}
           approval={part.approval}
           onApprove={addToolApprovalResponse}
-          isLive={isLast}
-          isActive={isActive}
+          isActive={running}
+          className={inActivity ? undefined : "my-1"}
         />
       );
     }
-
     if (isDataPart(part)) {
-      Logger.debug("Found data part", part);
       return null;
     }
 
@@ -114,26 +106,30 @@ export const renderUIMessage = ({
         if (part.text.includes("<marimo-")) {
           return (
             <React.Fragment key={index}>
-              {renderHTML({
-                html: part.text,
-              })}
+              {renderHTML({ html: part.text })}
             </React.Fragment>
           );
         }
-        return <ChatMarkdown key={index} content={part.text} />;
+        return (
+          <div key={index} className={inActivity ? "py-1.5" : undefined}>
+            <ChatMarkdown content={part.text} />
+          </div>
+        );
       case "reasoning":
         return (
-          <ReasoningAccordion
+          <div
             key={index}
-            reasoning={part.text}
-            index={index}
-            isStreaming={
-              isStreamingReasoning &&
-              isLast &&
-              // If there are multiple reasoning parts, only stream the last one
-              index === (message.parts.length || 0) - 1
-            }
-          />
+            className="flex gap-1.5 py-1.5 text-muted-foreground"
+          >
+            <BrainIcon
+              aria-hidden={true}
+              className="-ml-px mt-1 size-3 shrink-0"
+            />
+            <div className="min-w-0 italic [&_p]:my-1 [&_p:first-child]:mt-0">
+              <span className="sr-only">Reasoning: </span>
+              <ChatMarkdown content={part.text} />
+            </div>
+          </div>
         );
       case "file":
         return <AttachmentRenderer attachment={part} key={index} />;
@@ -148,23 +144,6 @@ export const renderUIMessage = ({
             }}
           />
         );
-      case "dynamic-tool":
-        return (
-          <ToolCallView
-            key={index}
-            toolName={part.toolName}
-            result={part.output}
-            errorText={
-              part.state === "output-error" ? part.errorText : undefined
-            }
-            state={part.state}
-            input={part.input}
-            approval={part.approval}
-            onApprove={addToolApprovalResponse}
-            isLive={isLast}
-            isActive={isActive}
-          />
-        );
       case "source-document":
         return (
           <SourceChip
@@ -199,12 +178,62 @@ export const renderUIMessage = ({
   }
 };
 
-function isToolPart(part: UIMessage["parts"][number]): part is ToolUIPart {
+export function getAnswerText(message: UIMessage): string {
+  const textOf = (parts: IndexedPart[]) =>
+    parts.flatMap(({ part }) => (part.type === "text" ? [part.text] : []));
+  const { work, rest } = splitTurn(message.parts);
+  const answer = textOf(rest);
+  return (answer.length > 0 ? answer : textOf(work)).join("\n");
+}
+
+function splitTurn(parts: UIMessage["parts"]): {
+  work: IndexedPart[];
+  rest: IndexedPart[];
+} {
+  const visible = parts
+    .map((part, index) => ({ part, index }))
+    .filter(({ part }) => isVisible(part));
+  const lastWork = visible.findLastIndex(
+    ({ part }) =>
+      part.type === "reasoning" ||
+      (isAnyToolPart(part) && !isPendingApproval(part)),
+  );
+  return {
+    work: visible.filter(
+      ({ part }, i) => i <= lastWork && !isPendingApproval(part),
+    ),
+    rest: visible.filter(
+      ({ part }, i) => i > lastWork || isPendingApproval(part),
+    ),
+  };
+}
+
+function isPendingApproval(part: Part): boolean {
+  return isAnyToolPart(part) && part.state === "approval-requested";
+}
+
+function isVisible(part: Part): boolean {
+  switch (part.type) {
+    case "step-start":
+      return false;
+    case "text":
+      return part.text.trim().length > 0;
+    case "reasoning":
+      // Some models (e.g. OpenAI's o-series) emit empty reasoning parts.
+      return part.text.trim().length > 0;
+    default:
+      return !isDataPart(part);
+  }
+}
+
+function isAnyToolPart(part: Part): part is ToolUIPart | DynamicToolUIPart {
+  return isToolPart(part) || part.type === "dynamic-tool";
+}
+
+function isToolPart(part: Part): part is ToolUIPart {
   return part.type.startsWith("tool-");
 }
 
-function isDataPart(
-  part: UIMessage["parts"][number],
-): part is DataUIPart<Record<string, unknown>> {
+function isDataPart(part: Part): part is DataUIPart<Record<string, unknown>> {
   return part.type.startsWith("data-");
 }

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

@@ -20,7 +20,6 @@ import {
 import { memo, useEffect, useMemo, useRef, useState } from "react";
 import useEvent from "react-use-event-hook";
 import { Button } from "@/components/ui/button";
-import { LoadingEllipsis } from "@/components/icons/loading-ellipsis";
 import { replaceMessagesInChat } from "@/core/ai/chat-utils";
 import { useAIConfigActions } from "@/core/ai/config";
 import { AI_SDK_UI_THROTTLE_MS } from "@/core/ai/constants";
@@ -43,7 +42,6 @@ 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";
@@ -52,6 +50,7 @@ import {
   usePairWithAgentModal,
 } from "../editor/actions/pair-with-agent";
 import { PromptInput } from "../editor/ai/prompt-input";
+import { PendingAICells } from "../editor/chrome/wrapper/pending-ai-cells";
 import {
   addContextCompletion,
   isContextAttachment,
@@ -76,7 +75,8 @@ 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";
@@ -87,7 +87,6 @@ import {
   handleToolCall,
   hasPendingToolCalls,
   hasUnresolvedToolCalls,
-  isLastMessageReasoning,
   PROVIDERS_THAT_SUPPORT_ATTACHMENTS,
   type QueuedUserMessage,
   shouldFlushQueue,
@@ -98,7 +97,9 @@ import { getCodes } from "@/core/codemirror/copilot/getCodes";
 import { focusInputAndMoveToEnd } from "@/core/codemirror/utils";
 import ScrollToBottomButton from "./acp/scroll-to-bottom-button";
 import { useChatControllerId } from "./use-chat-controller-id";
-import { useChatGeneration } from "./use-chat-generation";
+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 = "ask";
@@ -190,9 +191,9 @@ interface ChatMessageProps {
   message: UIMessage;
   index: number;
   onEdit: (index: number, newValue: string) => void;
-  isStreamingReasoning: boolean;
   isLast: boolean;
   isActive: boolean;
+  turnStartedAt: number | null;
   addToolApprovalResponse?: ChatAddToolApproveResponseFunction;
 }
 
@@ -201,78 +202,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-[10px]">
-          {fileParts?.map((filePart, idx) => (
-            <AttachmentRenderer attachment={filePart} key={idx} />
-          ))}
-          <PromptInput
-            key={message.id}
-            value={content}
-            placeholder="Type your message. Use @ to reference variables and cells."
-            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 group-focus-within:opacity-100 transition-opacity">
-            <CopyClipboardIcon className="h-3 w-3" value={content || ""} />
+            <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>
     );
   },
@@ -483,6 +446,8 @@ const ChatPanelBody = () => {
     lastAbortReasonRef.current = null;
     setAbortNotice(null);
   });
+  const turn = useAssistantTurn();
+  const resetTurn = turn.reset;
   const runtimeManager = useRuntimeManager();
   const { invokeAiTool } = useRequestClient();
 
@@ -505,6 +470,7 @@ const ChatPanelBody = () => {
     error,
     status,
     regenerate,
+    setMessages,
     stop,
     addToolOutput,
     addToolApprovalResponse,
@@ -512,7 +478,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(
       {
@@ -559,7 +526,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,
@@ -604,6 +590,7 @@ const ChatPanelBody = () => {
 
   const sendUserMessage = useEvent((parts: ChatMessagePart[]) => {
     clearAbortState();
+    turn.start();
     sendMessage({ role: "user", parts });
   });
 
@@ -628,11 +615,43 @@ const ChatPanelBody = () => {
     },
   );
 
-  const isLoading = useChatGeneration(status);
+  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 (isLoading || hasQueuedRef.current) {
+    if (isTurnActive || hasQueuedRef.current) {
       enqueueUserMessage(parts);
     } else {
       sendUserMessage(parts);
@@ -665,10 +684,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);
@@ -676,6 +691,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;
@@ -730,6 +750,7 @@ const ChatPanelBody = () => {
       const { contextPart, attachments } =
         await resolveChatContext(initialMessage);
 
+      turn.start();
       sendMessage({
         role: "user",
         parts: [
@@ -787,6 +808,7 @@ const ChatPanelBody = () => {
       const { contextPart, attachments } = await resolveChatContext(newValue);
 
       const messageId = editedMessage.id;
+      turn.start();
       sendMessage({
         messageId: messageId, // replace the message
         role: "user",
@@ -825,6 +847,7 @@ const ChatPanelBody = () => {
 
   const handleReload = () => {
     clearAbortState();
+    turn.start();
     regenerate();
   };
 
@@ -889,8 +912,8 @@ const ChatPanelBody = () => {
       inputRef={newThreadInputRef}
       setInput={setNewThreadInput}
       onSubmit={handleNewThreadSubmit}
-      isLoading={isLoading}
-      onStop={stop}
+      isLoading={isTurnActive}
+      onStop={stopTurn}
       fileInputRef={fileInputRef}
       onAddFiles={addFiles}
       onClose={handleOnCloseThread}
@@ -901,8 +924,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}
@@ -946,18 +969,35 @@ const ChatPanelBody = () => {
           ref={scrollContainerRef}
           onScroll={handleScroll}
         >
-          {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} />
@@ -990,30 +1030,7 @@ const ChatPanelBody = () => {
         />
       </div>
 
-      {isLoading && (
-        <div className="flex shrink-0 items-center gap-2 px-3 pt-2 text-xs text-muted-foreground">
-          <div className="flex items-center gap-1.5">
-            <div aria-hidden={true}>
-              <LoadingEllipsis
-                size={3}
-                gap="gap-0.5"
-                className="motion-reduce:[&_svg]:animate-none"
-              />
-            </div>
-            <output>Working</output>
-          </div>
-          <span aria-hidden={true} className="h-3 w-px bg-border" />
-          <Button
-            variant="text"
-            size="xs"
-            className="chat-control"
-            onClick={stop}
-          >
-            Stop generation
-          </Button>
-        </div>
-      )}
-
+      {!isTurnActive && <PendingAICells />}
       <div className="shrink-0 p-3">{composer}</div>
     </div>
   );