feat: refine token usage display modes (#2329)

* feat: refine token usage display modes * docs: clarify token usage accounting semantics * fix: avoid duplicate subtask debug keys * style: format token usage tests * chore: address token attribution review feedback * Update test_token_usage_middleware.py * Update test_token_usage_middleware.py * chore: simplify token attribution fallback * fix token usage metadata follow-up handling --------- Co-authored-by: Willem Jiang <willem.jiang@gmail.com>
2026-05-24 00:45:57 +00:00 · 2026-05-04 09:56:16 +08:00
parent 82e7936d36
commit d02f762ab0
20 changed files with 2346 additions and 222 deletions
@@ -1,6 +1,7 @@
+import type { Message } from "@langchain/langgraph-sdk";
 import type { BaseStream } from "@langchain/langgraph-sdk/react";
 import { ChevronUpIcon, Loader2Icon } from "lucide-react";
-import { useCallback, useEffect, useRef } from "react";
+import { useCallback, useEffect, useMemo, useRef } from "react";

 import {
  Conversation,
@@ -8,15 +9,20 @@ import {
 } from "@/components/ai-elements/conversation";
 import { Button } from "@/components/ui/button";
 import { useI18n } from "@/core/i18n/hooks";
+import {
+  buildTokenDebugSteps,
+  type TokenUsageInlineMode,
+} from "@/core/messages/usage-model";
 import {
  extractContentFromMessage,
  extractPresentFilesFromMessage,
+  extractReasoningContentFromMessage,
  extractTextFromMessage,
-  groupMessages,
+  getAssistantTurnUsageMessages,
+  getMessageGroups,
  hasContent,
  hasPresentFiles,
  hasReasoning,
-  hasToolCalls,
 } from "@/core/messages/utils";
 import { useRehypeSplitWordsIntoSpans } from "@/core/rehype";
 import type { Subtask } from "@/core/tasks";
@@ -25,12 +31,16 @@ import type { AgentThreadState } from "@/core/threads";
 import { cn } from "@/lib/utils";

 import { ArtifactFileList } from "../artifacts/artifact-file-list";
+import { CopyButton } from "../copy-button";
 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 {
+  MessageTokenUsageDebugList,
+  MessageTokenUsageList,
+} from "./message-token-usage";
 import { MessageListSkeleton } from "./skeleton";
 import { SubtaskCard } from "./subtask-card";

@@ -149,7 +159,7 @@ export function MessageList({
  threadId,
  thread,
  paddingBottom = MESSAGE_LIST_DEFAULT_PADDING_BOTTOM,
-  tokenUsageEnabled = false,
+  tokenUsageInlineMode = "off",
  hasMoreHistory,
  loadMoreHistory,
  isHistoryLoading,
@@ -158,7 +168,7 @@ export function MessageList({
  threadId: string;
  thread: BaseStream<AgentThreadState>;
  paddingBottom?: number;
-  tokenUsageEnabled?: boolean;
+  tokenUsageInlineMode?: TokenUsageInlineMode;
  hasMoreHistory?: boolean;
  loadMoreHistory?: () => void;
  isHistoryLoading?: boolean;
@@ -167,10 +177,85 @@ export function MessageList({
  const rehypePlugins = useRehypeSplitWordsIntoSpans(thread.isLoading);
  const updateSubtask = useUpdateSubtask();
  const messages = thread.messages;
+  const groupedMessages = getMessageGroups(messages);
+  const turnUsageMessagesByGroupIndex =
+    getAssistantTurnUsageMessages(groupedMessages);
+  const tokenDebugSteps = useMemo(
+    () => buildTokenDebugSteps(messages, t),
+    [messages, t],
+  );
+
+  const renderAssistantCopyButton = useCallback((messages: Message[]) => {
+    const clipboardData = [...messages]
+      .reverse()
+      .filter((message) => message.type === "ai")
+      .map((message) => {
+        const content = extractContentFromMessage(message);
+        return content ?? extractReasoningContentFromMessage(message) ?? "";
+      })
+      .find((content) => content.length > 0);
+
+    if (!clipboardData) {
+      return null;
+    }
+
+    return (
+      <div className="mt-2 flex justify-start opacity-0 transition-opacity delay-200 duration-300 group-hover/assistant-turn:opacity-100">
+        <CopyButton clipboardData={clipboardData} />
+      </div>
+    );
+  }, []);
+
+  const renderTokenUsage = useCallback(
+    ({
+      messages,
+      turnUsageMessages,
+      inlineDebug = true,
+      debugMessageIds,
+    }: {
+      messages: Message[];
+      turnUsageMessages?: Message[] | null;
+      inlineDebug?: boolean;
+      debugMessageIds?: string[];
+    }) => {
+      if (tokenUsageInlineMode === "per_turn") {
+        return (
+          <MessageTokenUsageList
+            enabled={true}
+            isLoading={thread.isLoading}
+            messages={turnUsageMessages ?? []}
+          />
+        );
+      }
+
+      if (tokenUsageInlineMode === "step_debug" && inlineDebug) {
+        const messageIds = new Set(
+          debugMessageIds ??
+            messages
+              .filter((message) => message.type === "ai")
+              .map((message) => message.id)
+              .filter((id): id is string => typeof id === "string"),
+        );
+        return (
+          <MessageTokenUsageDebugList
+            enabled={true}
+            isLoading={thread.isLoading}
+            steps={tokenDebugSteps.filter((step) =>
+              messageIds.has(step.messageId),
+            )}
+          />
+        );
+      }
+
+      return null;
+    },
+    [thread.isLoading, tokenDebugSteps, tokenUsageInlineMode],
+  );

  if (thread.isThreadLoading && messages.length === 0) {
    return <MessageListSkeleton />;
  }
+
  return (
    <Conversation
      className={cn("flex size-full flex-col justify-center", className)}
@@ -181,19 +266,37 @@ export function MessageList({
          hasMore={hasMoreHistory}
          loadMore={loadMoreHistory}
        />
-        {groupMessages(messages, (group) => {
+        {groupedMessages.map((group, groupIndex) => {
+          const turnUsageMessages = turnUsageMessagesByGroupIndex[groupIndex];
+
          if (group.type === "human" || group.type === "assistant") {
-            return group.messages.map((msg) => {
-              return (
-                <MessageListItem
-                  key={`${group.id}/${msg.id}`}
-                  threadId={threadId}
-                  message={msg}
-                  isLoading={thread.isLoading}
-                  tokenUsageEnabled={tokenUsageEnabled}
-                />
-              );
-            });
+            return (
+              <div
+                key={group.id}
+                className={cn(
+                  "w-full",
+                  group.type === "assistant" && "group/assistant-turn",
+                )}
+              >
+                {group.messages.map((msg) => {
+                  return (
+                    <MessageListItem
+                      key={`${group.id}/${msg.id}`}
+                      message={msg}
+                      isLoading={thread.isLoading}
+                      threadId={threadId}
+                      showCopyButton={group.type !== "assistant"}
+                    />
+                  );
+                })}
+                {renderTokenUsage({
+                  messages: group.messages,
+                  turnUsageMessages,
+                })}
+                {group.type === "assistant" &&
+                  renderAssistantCopyButton(group.messages)}
+              </div>
+            );
          } else if (group.type === "assistant:clarification") {
            const message = group.messages[0];
            if (message && hasContent(message)) {
@@ -204,11 +307,10 @@ export function MessageList({
                    isLoading={thread.isLoading}
                    rehypePlugins={rehypePlugins}
                  />
-                  <MessageTokenUsageList
-                    enabled={tokenUsageEnabled}
-                    isLoading={thread.isLoading}
-                    messages={group.messages}
-                  />
+                  {renderTokenUsage({
+                    messages: group.messages,
+                    turnUsageMessages,
+                  })}
                </div>
              );
            }
@@ -232,11 +334,10 @@ export function MessageList({
                  />
                )}
                <ArtifactFileList files={files} threadId={threadId} />
-                <MessageTokenUsageList
-                  enabled={tokenUsageEnabled}
-                  isLoading={thread.isLoading}
-                  messages={group.messages}
-                />
+                {renderTokenUsage({
+                  messages: group.messages,
+                  turnUsageMessages,
+                })}
              </div>
            );
          } else if (group.type === "assistant:subagent") {
@@ -289,7 +390,19 @@ export function MessageList({
                }
              }
            }
+
            const results: React.ReactNode[] = [];
+            const subagentDebugMessageIds: string[] = [];
+            if (tasks.size > 0) {
+              results.push(
+                <div
+                  key="subtask-count"
+                  className="text-muted-foreground pt-2 text-sm font-normal"
+                >
+                  {t.subtasks.executing(tasks.size)}
+                </div>,
+              );
+            }
            for (const message of group.messages.filter(
              (message) => message.type === "ai",
            )) {
@@ -299,17 +412,17 @@ export function MessageList({
                    key={"thinking-group-" + message.id}
                    messages={[message]}
                    isLoading={thread.isLoading}
+                    tokenDebugSteps={tokenDebugSteps.filter(
+                      (step) => step.messageId === message.id,
+                    )}
+                    showTokenDebugSummaries={
+                      tokenUsageInlineMode === "step_debug"
+                    }
                  />,
                );
+              } else if (message.id) {
+                subagentDebugMessageIds.push(message.id);
              }
-              results.push(
-                <div
-                  key="subtask-count"
-                  className="text-muted-foreground font-norma pt-2 text-sm"
-                >
-                  {t.subtasks.executing(tasks.size)}
-                </div>,
-              );
              const taskIds = message.tool_calls
                ?.filter((toolCall) => toolCall.name === "task")
                .map((toolCall) => toolCall.id);
@@ -329,30 +442,31 @@ export function MessageList({
                className="relative z-1 flex flex-col gap-2"
              >
                {results}
-                <MessageTokenUsageList
-                  enabled={tokenUsageEnabled}
-                  isLoading={thread.isLoading}
-                  messages={group.messages}
-                />
+                {renderTokenUsage({
+                  messages: group.messages,
+                  turnUsageMessages,
+                  debugMessageIds: subagentDebugMessageIds,
+                })}
              </div>
            );
          }
-          const tokenUsageMessages = group.messages.filter(
-            (message) =>
-              message.type === "ai" &&
-              (hasToolCalls(message) ? true : !hasContent(message)),
-          );
          return (
            <div key={"group-" + group.id} className="w-full">
              <MessageGroup
                messages={group.messages}
                isLoading={thread.isLoading}
+                tokenDebugSteps={tokenDebugSteps.filter((step) =>
+                  group.messages.some(
+                    (message) => message.id === step.messageId,
+                  ),
+                )}
+                showTokenDebugSummaries={tokenUsageInlineMode === "step_debug"}
              />
-              <MessageTokenUsageList
-                enabled={tokenUsageEnabled}
-                isLoading={thread.isLoading}
-                messages={tokenUsageMessages}
-              />
+              {renderTokenUsage({
+                messages: group.messages,
+                turnUsageMessages,
+                inlineDebug: false,
+              })}
            </div>
          );
        })}