feat(copilot): collapse repeated tool calls and fix stream stuck on completion

autogpt_platform/backend/backend/copilot/sdk/service.py

@@ -1196,18 +1196,20 @@ async def _next_msg() -> Any:
                 )
 
             if raw_transcript:
-                # Shield the upload from generator cancellation so a
-                # client disconnect / page refresh doesn't lose the
-                # transcript.  The upload must finish even if the SSE
-                # connection is torn down.
-                await asyncio.shield(
+                # Fire-and-forget: upload in background so the generator
+                # can exit promptly and mark_session_completed() runs
+                # without delay.  Blocking here kept the SSE stream alive
+                # with only heartbeats for up to 30s (the upload timeout).
+                task = asyncio.create_task(
                     _try_upload_transcript(
                         user_id,
                         session_id,
                         raw_transcript,
                         message_count=len(session.messages),
                     )
                 )
+                _background_tasks.add(task)
+                task.add_done_callback(_background_tasks.discard)
 
         logger.info(
             "[SDK] [%s] Stream completed successfully with %d messages",
@@ -1287,29 +1289,27 @@ async def _next_msg() -> Any:
                 )
 
         # --- Upload transcript for next-turn --resume ---
-        # This MUST run in finally so the transcript is uploaded even when
-        # the streaming loop raises an exception.  The CLI uses
-        # appendFileSync, so whatever was written before the error/SIGTERM
-        # is safely on disk and still useful for the next turn.
+        # Fire-and-forget so the generator exits promptly and
+        # mark_session_completed() in the processor runs without delay.
+        # Previously this awaited the upload (up to 30s timeout), keeping
+        # the SSE stream alive with only heartbeats.
         if config.claude_agent_use_resume and user_id:
             try:
-                # Prefer content captured in the Stop hook (read before
-                # cleanup removes the file).  Fall back to the resume
-                # file when the stop hook didn't fire (e.g. error before
-                # completion) so we don't lose the prior transcript.
                 raw_transcript = captured_transcript.raw_content or None
                 if not raw_transcript and use_resume and resume_file:
                     raw_transcript = read_transcript_file(resume_file)
 
                 if raw_transcript and session is not None:
-                    await asyncio.shield(
+                    task = asyncio.create_task(
                         _try_upload_transcript(
                             user_id,
                             session_id,
                             raw_transcript,
                             message_count=len(session.messages),
                         )
                     )
+                    _background_tasks.add(task)
+                    task.add_done_callback(_background_tasks.discard)
                 else:
                     logger.warning(f"[SDK] No transcript to upload for {session_id}")
             except Exception as upload_err:

autogpt_platform/frontend/src/app/(platform)/copilot/components/ChatContainer/ChatContainer.tsx

@@ -64,6 +64,7 @@ export const ChatContainer = ({
                 error={error}
                 isLoading={isLoadingSession}
                 headerSlot={headerSlot}
+                sessionID={sessionId}
               />
               <motion.div
                 initial={{ opacity: 0 }}

autogpt_platform/frontend/src/app/(platform)/copilot/components/ChatMessagesContainer/ChatMessagesContainer.tsx

@@ -3,26 +3,152 @@ import {
   ConversationContent,
   ConversationScrollButton,
 } from "@/components/ai-elements/conversation";
-import {
-  Message,
-  MessageActions,
-  MessageContent,
-} from "@/components/ai-elements/message";
+import { Message, MessageContent } from "@/components/ai-elements/message";
 import { LoadingSpinner } from "@/components/atoms/LoadingSpinner/LoadingSpinner";
-import { FileUIPart, UIDataTypes, UIMessage, UITools } from "ai";
+import { FileUIPart, ToolUIPart, UIDataTypes, UIMessage, UITools } from "ai";
+import { AssistantMessageActions } from "./components/AssistantMessageActions";
+import { CollapsedToolGroup } from "./components/CollapsedToolGroup";
 import { MessageAttachments } from "./components/MessageAttachments";
 import { MessagePartRenderer } from "./components/MessagePartRenderer";
+import { ReasoningCollapse } from "./components/ReasoningCollapse";
 import { ThinkingIndicator } from "./components/ThinkingIndicator";
-import { CopyButton } from "./components/CopyButton";
-import { TTSButton } from "./components/TTSButton";
-import { parseSpecialMarkers } from "./helpers";
+
+type MessagePart = UIMessage<unknown, UIDataTypes, UITools>["parts"][number];
 
 interface Props {
   messages: UIMessage<unknown, UIDataTypes, UITools>[];
   status: string;
   error: Error | undefined;
   isLoading: boolean;
   headerSlot?: React.ReactNode;
+  sessionID?: string | null;
+}
+
+function isCompletedToolPart(part: MessagePart): part is ToolUIPart {
+  return (
+    part.type.startsWith("tool-") &&
+    "state" in part &&
+    (part.state === "output-available" || part.state === "output-error")
+  );
+}
+
+type RenderSegment =
+  | { kind: "part"; part: MessagePart; index: number }
+  | { kind: "collapsed-group"; parts: ToolUIPart[] };
+
+// Tool types that have custom renderers and should NOT be collapsed
+const CUSTOM_TOOL_TYPES = new Set([
+  "tool-find_block",
+  "tool-find_agent",
+  "tool-find_library_agent",
+  "tool-search_docs",
+  "tool-get_doc_page",
+  "tool-run_block",
+  "tool-run_mcp_tool",
+  "tool-run_agent",
+  "tool-schedule_agent",
+  "tool-create_agent",
+  "tool-edit_agent",
+  "tool-view_agent_output",
+  "tool-search_feature_requests",
+  "tool-create_feature_request",
+]);
+
+/**
+ * Groups consecutive completed generic tool parts into collapsed segments.
+ * Non-generic tools (those with custom renderers) and active/streaming tools
+ * are left as individual parts.
+ */
+function buildRenderSegments(parts: MessagePart[]): RenderSegment[] {
+  const segments: RenderSegment[] = [];
+  let pendingGroup: ToolUIPart[] | null = null;
+
+  function flushGroup() {
+    if (!pendingGroup) return;
+    if (pendingGroup.length >= 2) {
+      segments.push({ kind: "collapsed-group", parts: pendingGroup });
+    } else {
+      for (const p of pendingGroup) {
+        const idx = parts.indexOf(p);
+        segments.push({ kind: "part", part: p, index: idx });
+      }
+    }
+    pendingGroup = null;
+  }
+
+  parts.forEach((part, i) => {
+    const isGenericCompletedTool =
+      isCompletedToolPart(part) && !CUSTOM_TOOL_TYPES.has(part.type);
+
+    if (isGenericCompletedTool) {
+      if (!pendingGroup) pendingGroup = [];
+      pendingGroup.push(part as ToolUIPart);
+    } else {
+      flushGroup();
+      segments.push({ kind: "part", part, index: i });
+    }
+  });
+
+  flushGroup();
+  return segments;
+}
+
+/**
+ * For finalized assistant messages, split parts into "reasoning" (intermediate
+ * text + tools before the final response) and "response" (final text after the
+ * last tool). If there are no tools, everything is response.
+ */
+function splitReasoningAndResponse(parts: MessagePart[]): {
+  reasoning: MessagePart[];
+  response: MessagePart[];
+} {
+  // Find the index of the last tool part
+  let lastToolIndex = -1;
+  for (let i = parts.length - 1; i >= 0; i--) {
+    if (parts[i].type.startsWith("tool-")) {
+      lastToolIndex = i;
+      break;
+    }
+  }
+
+  // No tools → everything is response
+  if (lastToolIndex === -1) {
+    return { reasoning: [], response: parts };
+  }
+
+  // Check if there's any text after the last tool
+  const hasResponseAfterTools = parts
+    .slice(lastToolIndex + 1)
+    .some((p) => p.type === "text");
+
+  if (!hasResponseAfterTools) {
+    // No final text response → don't collapse anything
+    return { reasoning: [], response: parts };
+  }
+
+  return {
+    reasoning: parts.slice(0, lastToolIndex + 1),
+    response: parts.slice(lastToolIndex + 1),
+  };
+}
+
+function renderSegments(
+  segments: RenderSegment[],
+  messageID: string,
+): React.ReactNode[] {
+  return segments.map((seg, segIdx) => {
+    if (seg.kind === "collapsed-group") {
+      return <CollapsedToolGroup key={`group-${segIdx}`} parts={seg.parts} />;
+    }
+    return (
+      <MessagePartRenderer
+        key={`${messageID}-${seg.index}`}
+        part={seg.part}
+        messageID={messageID}
+        partIndex={seg.index}
+      />
+    );
+  });
 }
 
 export function ChatMessagesContainer({
@@ -31,6 +157,7 @@ export function ChatMessagesContainer({
   error,
   isLoading,
   headerSlot,
+  sessionID,
 }: Props) {
   const lastMessage = messages[messages.length - 1];
 
@@ -80,19 +207,37 @@ export function ChatMessagesContainer({
             messageIndex === messages.length - 1 &&
             message.role === "assistant";
 
-          const isAssistant = message.role === "assistant";
-
-          // Past assistant messages are always done; the last one
-          // is done only when the stream has finished.
-          const isAssistantDone =
-            isAssistant &&
-            (!isLastAssistant ||
-              (status !== "streaming" && status !== "submitted"));
+          const isCurrentlyStreaming =
+            isLastAssistant &&
+            (status === "streaming" || status === "submitted");
+          const nextMessage = messages[messageIndex + 1];
+          const isLastInTurn =
+            message.role === "assistant" &&
+            (!nextMessage || nextMessage.role === "user");
+          const showActions =
+            isLastInTurn &&
+            !isCurrentlyStreaming &&
+            message.parts.some((p) => p.type === "text");
 
           const fileParts = message.parts.filter(
             (p): p is FileUIPart => p.type === "file",
           );
 
+          // For finalized assistant messages, split into reasoning + response.
+          // During streaming, show everything normally with tool collapsing.
+          const isFinalized =
+            message.role === "assistant" && !isCurrentlyStreaming;
+          const { reasoning, response } = isFinalized
+            ? splitReasoningAndResponse(message.parts)
+            : { reasoning: [] as MessagePart[], response: message.parts };
+          const hasReasoning = reasoning.length > 0;
+
+          const responseSegments =
+            message.role === "assistant" ? buildRenderSegments(response) : null;
+          const reasoningSegments = hasReasoning
+            ? buildRenderSegments(reasoning)
+            : null;
+
           return (
             <Message from={message.role} key={message.id}>
               <MessageContent
@@ -102,14 +247,21 @@ export function ChatMessagesContainer({
                   "group-[.is-assistant]:bg-transparent group-[.is-assistant]:text-slate-900"
                 }
               >
-                {message.parts.map((part, i) => (
-                  <MessagePartRenderer
-                    key={`${message.id}-${i}`}
-                    part={part}
-                    messageID={message.id}
-                    partIndex={i}
-                  />
-                ))}
+                {hasReasoning && reasoningSegments && (
+                  <ReasoningCollapse>
+                    {renderSegments(reasoningSegments, message.id)}
+                  </ReasoningCollapse>
+                )}
+                {responseSegments
+                  ? renderSegments(responseSegments, message.id)
+                  : message.parts.map((part, i) => (
+                      <MessagePartRenderer
+                        key={`${message.id}-${i}`}
+                        part={part}
+                        messageID={message.id}
+                        partIndex={i}
+                      />
+                    ))}
                 {isLastAssistant && showThinking && (
                   <ThinkingIndicator active={showThinking} />
                 )}
@@ -120,30 +272,12 @@ export function ChatMessagesContainer({
                   isUser={message.role === "user"}
                 />
               )}
-              {isAssistantDone &&
-                (() => {
-                  const textParts = message.parts.filter(
-                    (p): p is Extract<typeof p, { type: "text" }> =>
-                      p.type === "text",
-                  );
-
-                  // Hide actions when the message ended with an error or cancellation
-                  const lastTextPart = textParts[textParts.length - 1];
-                  if (lastTextPart) {
-                    const { markerType } = parseSpecialMarkers(
-                      lastTextPart.text,
-                    );
-                    if (markerType === "error") return null;
-                  }
-
-                  const textContent = textParts.map((p) => p.text).join("\n");
-                  return (
-                    <MessageActions>
-                      <CopyButton text={textContent} />
-                      <TTSButton text={textContent} />
-                    </MessageActions>
-                  );
-                })()}
+              {showActions && (
+                <AssistantMessageActions
+                  message={message}
+                  sessionID={sessionID ?? null}
+                />
+              )}
             </Message>
           );
         })}