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

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

@@ -827,6 +827,7 @@ async def stream_chat_completion_sdk(
     use_resume = False
     resume_file: str | None = None
     captured_transcript = CapturedTranscript()
+    transcript_upload_scheduled = False
     sdk_cwd = ""
 
     # Acquire stream lock to prevent concurrent streams to the same session
@@ -1436,18 +1437,21 @@ 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)
+                transcript_upload_scheduled = True
 
         logger.info(
             "[SDK] [%s] Stream completed successfully with %d messages",
@@ -1527,29 +1531,31 @@ 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.
-        if config.claude_agent_use_resume and user_id:
+        # 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
+            and not transcript_upload_scheduled
+        ):
             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/ChatMessagesContainer/ChatMessagesContainer.tsx

@@ -5,15 +5,19 @@ import {
 } from "@/components/ai-elements/conversation";
 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 { TOOL_PART_PREFIX } from "../JobStatsBar/constants";
 import { TurnStatsBar } from "../JobStatsBar/TurnStatsBar";
 import { parseSpecialMarkers } from "./helpers";
 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";
 
+type MessagePart = UIMessage<unknown, UIDataTypes, UITools>["parts"][number];
+
 interface Props {
   messages: UIMessage<unknown, UIDataTypes, UITools>[];
   status: string;
@@ -23,6 +27,139 @@ interface Props {
   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[],
+  baseIndex = 0,
+): RenderSegment[] {
+  const segments: RenderSegment[] = [];
+  let pendingGroup: Array<{ part: ToolUIPart; index: number }> | null = null;
+
+  function flushGroup() {
+    if (!pendingGroup) return;
+    if (pendingGroup.length >= 2) {
+      segments.push({
+        kind: "collapsed-group",
+        parts: pendingGroup.map((p) => p.part),
+      });
+    } else {
+      for (const p of pendingGroup) {
+        segments.push({ kind: "part", part: p.part, index: p.index });
+      }
+    }
+    pendingGroup = null;
+  }
+
+  parts.forEach((part, i) => {
+    const absoluteIndex = baseIndex + i;
+    const isGenericCompletedTool =
+      isCompletedToolPart(part) && !CUSTOM_TOOL_TYPES.has(part.type);
+
+    if (isGenericCompletedTool) {
+      if (!pendingGroup) pendingGroup = [];
+      pendingGroup.push({ part: part as ToolUIPart, index: absoluteIndex });
+    } else {
+      flushGroup();
+      segments.push({ kind: "part", part, index: absoluteIndex });
+    }
+  });
+
+  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}
+      />
+    );
+  });
+}
+
 /** Collect all messages belonging to a turn: the user message + every
  *  assistant message up to (but not including) the next user message. */
 function getTurnMessages(
@@ -119,6 +256,24 @@ export function ChatMessagesContainer({
             (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 responseStartIndex = message.parts.length - response.length;
+          const responseSegments =
+            message.role === "assistant"
+              ? buildRenderSegments(response, responseStartIndex)
+              : null;
+          const reasoningSegments = hasReasoning
+            ? buildRenderSegments(reasoning, 0)
+            : null;
+
           return (
             <Message from={message.role} key={message.id}>
               <MessageContent
@@ -128,14 +283,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}
+                      />
+                    ))}
                 {isLastInTurn && !isCurrentlyStreaming && (
                   <TurnStatsBar
                     turnMessages={getTurnMessages(messages, messageIndex)}