feat: enhance message display

frontend/src/components/workspace/messages/message-list-item.tsx

@@ -1,5 +1,5 @@
 import type { Message } from "@langchain/langgraph-sdk";
-import { memo } from "react";
+import { memo, useMemo } from "react";
 
 import {
   Message as AIElementMessage,
@@ -9,25 +9,20 @@ import {
 } from "@/components/ai-elements/message";
 import {
   extractContentFromMessage,
-  hasReasoning,
-  hasToolCalls,
+  extractReasoningContentFromMessage,
 } from "@/core/messages/utils";
 import { useRehypeSplitWordsIntoSpans } from "@/core/rehype";
 import { cn } from "@/lib/utils";
 
 import { CopyButton } from "../copy-button";
 
-import { MessageGroup } from "./message-group";
-
 export function MessageListItem({
   className,
   message,
-  messagesInGroup,
   isLoading,
 }: {
   className?: string;
   message: Message;
-  messagesInGroup: Message[];
   isLoading?: boolean;
 }) {
   return (
@@ -38,7 +33,6 @@ export function MessageListItem({
       <MessageContent
         className={message.type === "human" ? "w-fit" : "w-full"}
         message={message}
-        messagesInGroup={messagesInGroup}
         isLoading={isLoading}
       />
       <MessageToolbar
@@ -49,7 +43,13 @@ export function MessageListItem({
         )}
       >
         <div className="flex gap-1">
-          <CopyButton clipboardData={extractContentFromMessage(message)} />
+          <CopyButton
+            clipboardData={
+              extractContentFromMessage(message)
+                ? extractContentFromMessage(message)
+                : (extractReasoningContentFromMessage(message) ?? "")
+            }
+          />
         </div>
       </MessageToolbar>
     </AIElementMessage>
@@ -59,26 +59,26 @@ export function MessageListItem({
 function MessageContent_({
   className,
   message,
-  messagesInGroup,
   isLoading = false,
 }: {
   className?: string;
   message: Message;
-  messagesInGroup: Message[];
   isLoading?: boolean;
 }) {
   const rehypePlugins = useRehypeSplitWordsIntoSpans(isLoading);
+  const content = useMemo(() => {
+    const reasoningContent = extractReasoningContentFromMessage(message);
+    const content = extractContentFromMessage(message);
+    if (!isLoading && reasoningContent && !content) {
+      return reasoningContent;
+    }
+    return content;
+  }, [isLoading, message]);
   return (
     <AIElementMessageContent className={className}>
-      {hasReasoning(message) && (
-        <MessageGroup messages={messagesInGroup} isLoading={isLoading} />
-      )}
       <AIElementMessageResponse rehypePlugins={rehypePlugins}>
-        {extractContentFromMessage(message)}
+        {content}
       </AIElementMessageResponse>
-      {hasToolCalls(message) && (
-        <MessageGroup messages={messagesInGroup} isLoading={isLoading} />
-      )}
     </AIElementMessageContent>
   );
 }

frontend/src/components/workspace/messages/message-list.tsx

@@ -7,7 +7,6 @@ import {
 import {
   extractPresentFilesFromMessage,
   groupMessages,
-  hasContent,
   hasPresentFiles,
 } from "@/core/messages/utils";
 import type { AgentThreadState } from "@/core/threads";
@@ -39,28 +38,26 @@ export function MessageList({
       <ConversationContent className="mx-auto w-full max-w-(--container-width-md) gap-10 pt-12">
         {groupMessages(
           thread.messages,
-          (groupedMessages) => {
-            if (groupedMessages[0] && hasContent(groupedMessages[0])) {
-              const message = groupedMessages[0];
+          (group) => {
+            if (group.type === "human" || group.type === "assistant") {
               return (
                 <MessageListItem
-                  key={message.id}
-                  message={message}
-                  messagesInGroup={groupedMessages}
+                  key={group.id}
+                  message={group.messages[0]!}
                   isLoading={thread.isLoading}
                 />
               );
             }
-            if (groupedMessages[0] && hasPresentFiles(groupedMessages[0])) {
+            if (group.type === "assistant:present-files") {
               const files = [];
-              for (const message of groupedMessages) {
+              for (const message of group.messages) {
                 if (hasPresentFiles(message)) {
                   files.push(...extractPresentFilesFromMessage(message));
                 }
               }
               return (
                 <ArtifactFileList
-                  key={groupedMessages[0].id}
+                  key={group.id}
                   files={files}
                   threadId={threadId}
                 />
@@ -68,8 +65,8 @@ export function MessageList({
             }
             return (
               <MessageGroup
-                key={groupedMessages[0]!.id}
-                messages={groupedMessages}
+                key={group.id}
+                messages={group.messages}
                 isLoading={thread.isLoading}
               />
             );

frontend/src/core/messages/utils.ts

@@ -1,58 +1,110 @@
 import type { Message } from "@langchain/langgraph-sdk";
 
+interface GenericMessageGroup<T = string> {
+  type: T;
+  id: string | undefined;
+  messages: Message[];
+}
+
+interface HumanMessageGroup extends GenericMessageGroup<"human"> {}
+
+interface AssistantProcessingGroup extends GenericMessageGroup<"assistant:processing"> {}
+
+interface AssistantMessageGroup extends GenericMessageGroup<"assistant"> {}
+
+interface AssistantPresentFilesGroup extends GenericMessageGroup<"assistant:present-files"> {}
+
+type MessageGroup =
+  | HumanMessageGroup
+  | AssistantProcessingGroup
+  | AssistantMessageGroup
+  | AssistantPresentFilesGroup;
+
 export function groupMessages<T>(
   messages: Message[],
-  mapper: (groupedMessages: Message[]) => T,
+  mapper: (group: MessageGroup) => T,
   isLoading = false,
 ): T[] {
   if (messages.length === 0) {
     return [];
   }
-  const groups: Message[][] = [];
-  let currentGroup: Message[] = [];
-  const yieldCurrentGroup = () => {
-    if (currentGroup.length > 0) {
-      groups.push(currentGroup);
-      currentGroup = [];
-    }
-  };
-  let messageIndex = 0;
+  const groups: MessageGroup[] = [];
+
   for (const message of messages) {
+    const lastGroup = groups[groups.length - 1];
     if (message.type === "human") {
-      // Human messages are always shown as a individual group
-      yieldCurrentGroup();
-      currentGroup.push(message);
-      yieldCurrentGroup();
+      groups.push({
+        id: message.id,
+        type: "human",
+        messages: [message],
+      });
     } else if (message.type === "tool") {
-      // Tool messages are always shown with the assistant messages that contains the tool calls
-      currentGroup.push(message);
-    } else if (message.type === "ai") {
       if (
-        hasToolCalls(message) ||
-        (extractTextFromMessage(message) === "" &&
-          extractReasoningContentFromMessage(message) !== "" &&
-          messageIndex === messages.length - 1 &&
-          isLoading)
+        lastGroup &&
+        lastGroup.type !== "human" &&
+        lastGroup.type !== "assistant"
       ) {
-        if (message.tool_calls?.[0]?.name === "present_files") {
-          // When `present_files` called, put them into an individual group
-          yieldCurrentGroup();
-          currentGroup.push(message);
-        } else {
-          // Assistant messages without any content are folded into the previous group
-          // Normally, these are tool calls (with or without thinking)
-          currentGroup.push(message);
-        }
+        lastGroup.messages.push(message);
       } else {
-        // Assistant messages with content (text or images) are shown as a group if they have content
-        // No matter whether it has tool calls or not
-        yieldCurrentGroup();
-        currentGroup.push(message);
+        throw new Error(
+          "Tool message must be matched with a previous assistant message with tool calls",
+        );
+      }
+    } else if (message.type === "ai") {
+      if (hasReasoning(message) || hasToolCalls(message)) {
+        if (hasPresentFiles(message)) {
+          groups.push({
+            id: message.id,
+            type: "assistant:present-files",
+            messages: [message],
+          });
+        } else {
+          if (lastGroup?.type !== "assistant:processing") {
+            groups.push({
+              id: message.id,
+              type: "assistant:processing",
+              messages: [],
+            });
+          }
+          const currentGroup = groups[groups.length - 1];
+          if (currentGroup?.type === "assistant:processing") {
+            currentGroup.messages.push(message);
+          } else {
+            throw new Error(
+              "Assistant message with reasoning or tool calls must be preceded by a processing group",
+            );
+          }
+        }
+      }
+      if (hasContent(message) && !hasToolCalls(message)) {
+        groups.push({
+          id: message.id,
+          type: "assistant",
+          messages: [message],
+        });
+      }
+    }
+  }
+
+  if (!isLoading) {
+    const lastGroup: MessageGroup | undefined = groups[groups.length - 1];
+    if (
+      lastGroup?.type === "assistant:processing" &&
+      lastGroup.messages.length > 0
+    ) {
+      const reasoningContent = extractReasoningContentFromMessage(
+        lastGroup.messages[lastGroup.messages.length - 1]!,
+      );
+      const content = extractContentFromMessage(
+        lastGroup.messages[lastGroup.messages.length - 1]!,
+      );
+      if (reasoningContent && !content) {
+        const group = groups.pop()!;
+        group.type = "assistant";
+        groups.push(group);
       }
     }
-    messageIndex++;
   }
-  yieldCurrentGroup();
 
   const resultsOfGroups: T[] = [];
   for (const group of groups) {