feat: show token usage per assistant response (#2270)

* feat: show token usage per assistant response

* fix: align client models response with token usage

* fix: address token usage review feedback

* docs: clarify token usage config example

---------

Co-authored-by: Willem Jiang <willem.jiang@gmail.com>

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

@@ -13,6 +13,7 @@ import {
   hasContent,
   hasPresentFiles,
   hasReasoning,
+  hasToolCalls,
 } from "@/core/messages/utils";
 import { useRehypeSplitWordsIntoSpans } from "@/core/rehype";
 import type { Subtask } from "@/core/tasks";
@@ -26,6 +27,7 @@ import { StreamingIndicator } from "../streaming-indicator";
 import { MarkdownContent } from "./markdown-content";
 import { MessageGroup } from "./message-group";
 import { MessageListItem } from "./message-list-item";
+import { MessageTokenUsageList } from "./message-token-usage";
 import { MessageListSkeleton } from "./skeleton";
 import { SubtaskCard } from "./subtask-card";
 
@@ -37,11 +39,13 @@ export function MessageList({
   threadId,
   thread,
   paddingBottom = MESSAGE_LIST_DEFAULT_PADDING_BOTTOM,
+  tokenUsageEnabled = false,
 }: {
   className?: string;
   threadId: string;
   thread: BaseStream<AgentThreadState>;
   paddingBottom?: number;
+  tokenUsageEnabled?: boolean;
 }) {
   const { t } = useI18n();
   const rehypePlugins = useRehypeSplitWordsIntoSpans(thread.isLoading);
@@ -64,6 +68,7 @@ export function MessageList({
                   message={msg}
                   isLoading={thread.isLoading}
                   threadId={threadId}
+                  tokenUsageEnabled={tokenUsageEnabled}
                 />
               );
             });
@@ -71,12 +76,18 @@ export function MessageList({
             const message = group.messages[0];
             if (message && hasContent(message)) {
               return (
-                <MarkdownContent
-                  key={group.id}
-                  content={extractContentFromMessage(message)}
-                  isLoading={thread.isLoading}
-                  rehypePlugins={rehypePlugins}
-                />
+                <div key={group.id} className="w-full">
+                  <MarkdownContent
+                    content={extractContentFromMessage(message)}
+                    isLoading={thread.isLoading}
+                    rehypePlugins={rehypePlugins}
+                  />
+                  <MessageTokenUsageList
+                    enabled={tokenUsageEnabled}
+                    isLoading={thread.isLoading}
+                    messages={group.messages}
+                  />
+                </div>
               );
             }
             return null;
@@ -99,6 +110,11 @@ export function MessageList({
                   />
                 )}
                 <ArtifactFileList files={files} threadId={threadId} />
+                <MessageTokenUsageList
+                  enabled={tokenUsageEnabled}
+                  isLoading={thread.isLoading}
+                  messages={group.messages}
+                />
               </div>
             );
           } else if (group.type === "assistant:subagent") {
@@ -191,15 +207,31 @@ export function MessageList({
                 className="relative z-1 flex flex-col gap-2"
               >
                 {results}
+                <MessageTokenUsageList
+                  enabled={tokenUsageEnabled}
+                  isLoading={thread.isLoading}
+                  messages={group.messages}
+                />
               </div>
             );
           }
+          const tokenUsageMessages = group.messages.filter(
+            (message) =>
+              message.type === "ai" &&
+              (hasToolCalls(message) ? true : !hasContent(message)),
+          );
           return (
-            <MessageGroup
-              key={"group-" + group.id}
-              messages={group.messages}
-              isLoading={thread.isLoading}
-            />
+            <div key={"group-" + group.id} className="w-full">
+              <MessageGroup
+                messages={group.messages}
+                isLoading={thread.isLoading}
+              />
+              <MessageTokenUsageList
+                enabled={tokenUsageEnabled}
+                isLoading={thread.isLoading}
+                messages={tokenUsageMessages}
+              />
+            </div>
           );
         })}
         {thread.isLoading && <StreamingIndicator className="my-4" />}