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/app/workspace/agents/[agent_name]/chats/[thread_id]/page.tsx

@@ -23,6 +23,7 @@ import { TokenUsageIndicator } from "@/components/workspace/token-usage-indicato
 import { Tooltip } from "@/components/workspace/tooltip";
 import { useAgent } from "@/core/agents";
 import { useI18n } from "@/core/i18n/hooks";
+import { useModels } from "@/core/models/hooks";
 import { useNotification } from "@/core/notification/hooks";
 import { useThreadSettings } from "@/core/settings";
 import { useThreadStream } from "@/core/threads/hooks";
@@ -44,6 +45,7 @@ export default function AgentChatPage() {
   const { threadId, setThreadId, isNewThread, setIsNewThread } =
     useThreadChat();
   const [settings, setSettings] = useThreadSettings(threadId);
+  const { tokenUsageEnabled } = useModels();
 
   const { showNotification } = useNotification();
   const [thread, sendMessage] = useThreadStream({
@@ -128,7 +130,10 @@ export default function AgentChatPage() {
                   <PlusSquare /> {t.agents.newChat}
                 </Button>
               </Tooltip>
-              <TokenUsageIndicator messages={thread.messages} />
+              <TokenUsageIndicator
+                enabled={tokenUsageEnabled}
+                messages={thread.messages}
+              />
               <ExportTrigger threadId={threadId} />
               <ArtifactTrigger />
             </div>
@@ -141,6 +146,7 @@ export default function AgentChatPage() {
                 threadId={threadId}
                 thread={thread}
                 paddingBottom={messageListPaddingBottom}
+                tokenUsageEnabled={tokenUsageEnabled}
               />
             </div>
 

frontend/src/app/workspace/chats/[thread_id]/page.tsx

@@ -22,6 +22,7 @@ import { TodoList } from "@/components/workspace/todo-list";
 import { TokenUsageIndicator } from "@/components/workspace/token-usage-indicator";
 import { Welcome } from "@/components/workspace/welcome";
 import { useI18n } from "@/core/i18n/hooks";
+import { useModels } from "@/core/models/hooks";
 import { useNotification } from "@/core/notification/hooks";
 import { useThreadSettings } from "@/core/settings";
 import { useThreadStream } from "@/core/threads/hooks";
@@ -36,6 +37,7 @@ export default function ChatPage() {
     useThreadChat();
   const [settings, setSettings] = useThreadSettings(threadId);
   const [mounted, setMounted] = useState(false);
+  const { tokenUsageEnabled } = useModels();
   useSpecificChatMode();
 
   useEffect(() => {
@@ -103,7 +105,10 @@ export default function ChatPage() {
               <ThreadTitle threadId={threadId} thread={thread} />
             </div>
             <div className="flex items-center gap-2">
-              <TokenUsageIndicator messages={thread.messages} />
+              <TokenUsageIndicator
+                enabled={tokenUsageEnabled}
+                messages={thread.messages}
+              />
               <ExportTrigger threadId={threadId} />
               <ArtifactTrigger />
             </div>
@@ -115,6 +120,7 @@ export default function ChatPage() {
                 threadId={threadId}
                 thread={thread}
                 paddingBottom={messageListPaddingBottom}
+                tokenUsageEnabled={tokenUsageEnabled}
               />
             </div>
             <div className="absolute right-0 bottom-0 left-0 z-30 flex justify-center px-4">

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

@@ -38,17 +38,20 @@ import { cn } from "@/lib/utils";
 import { CopyButton } from "../copy-button";
 
 import { MarkdownContent } from "./markdown-content";
+import { MessageTokenUsage } from "./message-token-usage";
 
 export function MessageListItem({
   className,
   message,
   isLoading,
   threadId,
+  tokenUsageEnabled = false,
 }: {
   className?: string;
   message: Message;
   isLoading?: boolean;
   threadId: string;
+  tokenUsageEnabled?: boolean;
 }) {
   const isHuman = message.type === "human";
   return (
@@ -61,6 +64,7 @@ export function MessageListItem({
         message={message}
         isLoading={isLoading}
         threadId={threadId}
+        tokenUsageEnabled={tokenUsageEnabled}
       />
       {!isLoading && (
         <MessageToolbar
@@ -119,11 +123,13 @@ function MessageContent_({
   message,
   isLoading = false,
   threadId,
+  tokenUsageEnabled = false,
 }: {
   className?: string;
   message: Message;
   isLoading?: boolean;
   threadId: string;
+  tokenUsageEnabled?: boolean;
 }) {
   const rehypePlugins = useRehypeSplitWordsIntoSpans(isLoading);
   const isHuman = message.type === "human";
@@ -201,6 +207,11 @@ function MessageContent_({
           <ReasoningTrigger />
           <ReasoningContent>{reasoningContent}</ReasoningContent>
         </Reasoning>
+        <MessageTokenUsage
+          enabled={tokenUsageEnabled}
+          isLoading={isLoading}
+          message={message}
+        />
       </AIElementMessageContent>
     );
   }
@@ -238,6 +249,11 @@ function MessageContent_({
         className="my-3"
         components={components}
       />
+      <MessageTokenUsage
+        enabled={tokenUsageEnabled}
+        isLoading={isLoading}
+        message={message}
+      />
     </AIElementMessageContent>
   );
 }

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" />}

frontend/src/components/workspace/messages/message-token-usage.tsx

@@ -0,0 +1,91 @@
+import type { Message } from "@langchain/langgraph-sdk";
+import { CoinsIcon } from "lucide-react";
+
+import { useI18n } from "@/core/i18n/hooks";
+import { formatTokenCount, getUsageMetadata } from "@/core/messages/usage";
+import { cn } from "@/lib/utils";
+
+export function MessageTokenUsage({
+  className,
+  enabled = false,
+  isLoading = false,
+  message,
+}: {
+  className?: string;
+  enabled?: boolean;
+  isLoading?: boolean;
+  message: Message;
+}) {
+  const { t } = useI18n();
+
+  if (!enabled || isLoading || message.type !== "ai") {
+    return null;
+  }
+
+  const usage = getUsageMetadata(message);
+
+  return (
+    <div
+      className={cn(
+        "text-muted-foreground border-border/60 mt-1 flex flex-wrap items-center gap-x-3 gap-y-1 border-t pt-2 text-[11px]",
+        className,
+      )}
+    >
+      <span className="inline-flex items-center gap-1 font-medium">
+        <CoinsIcon className="size-3" />
+        {t.tokenUsage.label}
+      </span>
+      {usage ? (
+        <>
+          <span>
+            {t.tokenUsage.input}: {formatTokenCount(usage.inputTokens)}
+          </span>
+          <span>
+            {t.tokenUsage.output}: {formatTokenCount(usage.outputTokens)}
+          </span>
+          <span className="font-medium">
+            {t.tokenUsage.total}: {formatTokenCount(usage.totalTokens)}
+          </span>
+        </>
+      ) : (
+        <span>{t.tokenUsage.unavailableShort}</span>
+      )}
+    </div>
+  );
+}
+
+export function MessageTokenUsageList({
+  className,
+  enabled = false,
+  isLoading = false,
+  messages,
+}: {
+  className?: string;
+  enabled?: boolean;
+  isLoading?: boolean;
+  messages: Message[];
+}) {
+  if (!enabled || isLoading) {
+    return null;
+  }
+
+  const aiMessages = messages.filter((message) => message.type === "ai");
+
+  if (aiMessages.length === 0) {
+    return null;
+  }
+
+  return (
+    <>
+      {aiMessages.map((message, index) => (
+        <MessageTokenUsage
+          className={className}
+          enabled={enabled}
+          isLoading={isLoading}
+          key={message.id ?? index}
+          message={message}
+        />
+      ))}
+    </>
+  );
+}

frontend/src/components/workspace/token-usage-indicator.tsx

@@ -15,18 +15,20 @@ import { cn } from "@/lib/utils";
 
 interface TokenUsageIndicatorProps {
   messages: Message[];
+  enabled?: boolean;
   className?: string;
 }
 
 export function TokenUsageIndicator({
   messages,
+  enabled = false,
   className,
 }: TokenUsageIndicatorProps) {
   const { t } = useI18n();
 
   const usage = useMemo(() => accumulateUsage(messages), [messages]);
 
-  if (!usage) {
+  if (!enabled) {
     return null;
   }
 
@@ -36,37 +38,49 @@ export function TokenUsageIndicator({
         <button
           type="button"
           className={cn(
-            "text-muted-foreground flex cursor-default items-center gap-1 text-xs",
+            "text-muted-foreground bg-background/70 flex cursor-default items-center gap-1.5 rounded-full border px-2 py-1 text-xs",
+            !usage && "opacity-60",
             className,
           )}
         >
           <CoinsIcon size={14} />
-          <span>{formatTokenCount(usage.totalTokens)}</span>
+          <span>{t.tokenUsage.label}</span>
+          <span className="font-mono">
+            {usage ? formatTokenCount(usage.totalTokens) : "-"}
+          </span>
         </button>
       </TooltipTrigger>
       <TooltipContent side="bottom" align="end">
         <div className="space-y-1 text-xs">
           <div className="font-medium">{t.tokenUsage.title}</div>
-          <div className="flex justify-between gap-4">
-            <span>{t.tokenUsage.input}</span>
-            <span className="font-mono">
-              {formatTokenCount(usage.inputTokens)}
-            </span>
-          </div>
-          <div className="flex justify-between gap-4">
-            <span>{t.tokenUsage.output}</span>
-            <span className="font-mono">
-              {formatTokenCount(usage.outputTokens)}
-            </span>
-          </div>
-          <div className="border-t pt-1">
-            <div className="flex justify-between gap-4">
-              <span>{t.tokenUsage.total}</span>
-              <span className="font-mono font-medium">
-                {formatTokenCount(usage.totalTokens)}
-              </span>
+          {usage ? (
+            <>
+              <div className="flex justify-between gap-4">
+                <span>{t.tokenUsage.input}</span>
+                <span className="font-mono">
+                  {formatTokenCount(usage.inputTokens)}
+                </span>
+              </div>
+              <div className="flex justify-between gap-4">
+                <span>{t.tokenUsage.output}</span>
+                <span className="font-mono">
+                  {formatTokenCount(usage.outputTokens)}
+                </span>
+              </div>
+              <div className="border-t pt-1">
+                <div className="flex justify-between gap-4">
+                  <span>{t.tokenUsage.total}</span>
+                  <span className="font-mono font-medium">
+                    {formatTokenCount(usage.totalTokens)}
+                  </span>
+                </div>
+              </div>
+            </>
+          ) : (
+            <div className="text-muted-foreground max-w-56">
+              {t.tokenUsage.unavailable}
             </div>
-          </div>
+          )}
         </div>
       </TooltipContent>
     </Tooltip>

frontend/src/core/i18n/locales/en-US.ts

@@ -298,9 +298,13 @@ export const enUS: Translations = {
   // Token Usage
   tokenUsage: {
     title: "Token Usage",
+    label: "Tokens",
     input: "Input",
     output: "Output",
     total: "Total",
+    unavailable:
+      "No token usage yet. Usage appears only after a successful model response when the provider returns usage_metadata.",
+    unavailableShort: "No usage returned",
   },
 
   // Shortcuts

frontend/src/core/i18n/locales/types.ts

@@ -229,9 +229,12 @@ export interface Translations {
   // Token Usage
   tokenUsage: {
     title: string;
+    label: string;
     input: string;
     output: string;
     total: string;
+    unavailable: string;
+    unavailableShort: string;
   };
 
   // Shortcuts

frontend/src/core/i18n/locales/zh-CN.ts

@@ -284,9 +284,13 @@ export const zhCN: Translations = {
   // Token Usage
   tokenUsage: {
     title: "Token 用量",
+    label: "Tokens",
     input: "输入",
     output: "输出",
     total: "总计",
+    unavailable:
+      "暂无 Token 用量。只有模型成功返回且供应商提供 usage_metadata 时才会显示。",
+    unavailableShort: "未返回用量",
   },
 
   // Shortcuts

frontend/src/core/messages/usage.ts

@@ -10,7 +10,7 @@ export interface TokenUsage {
  * Extract usage_metadata from an AI message if present.
  * The field is added by the backend (PR #1218) but not typed in the SDK.
  */
-function getUsageMetadata(message: Message): TokenUsage | null {
+export function getUsageMetadata(message: Message): TokenUsage | null {
   if (message.type !== "ai") {
     return null;
   }

frontend/src/core/models/api.ts

@@ -1,9 +1,12 @@
 import { getBackendBaseURL } from "../config";
 
-import type { Model } from "./types";
+import type { ModelsResponse } from "./types";
 
-export async function loadModels() {
+export async function loadModels(): Promise<ModelsResponse> {
   const res = await fetch(`${getBackendBaseURL()}/api/models`);
-  const { models } = (await res.json()) as { models: Model[] };
-  return models;
+  const data = (await res.json()) as Partial<ModelsResponse>;
+  return {
+    models: data.models ?? [],
+    token_usage: data.token_usage ?? { enabled: false },
+  };
 }

frontend/src/core/models/hooks.ts

@@ -9,5 +9,10 @@ export function useModels({ enabled = true }: { enabled?: boolean } = {}) {
     enabled,
     refetchOnWindowFocus: false,
   });
-  return { models: data ?? [], isLoading, error };
+  return {
+    models: data?.models ?? [],
+    tokenUsageEnabled: data?.token_usage.enabled ?? false,
+    isLoading,
+    error,
+  };
 }

frontend/src/core/models/types.ts

@@ -7,3 +7,12 @@ export interface Model {
   supports_thinking?: boolean;
   supports_reasoning_effort?: boolean;
 }
+
+export interface TokenUsageSettings {
+  enabled: boolean;
+}
+
+export interface ModelsResponse {
+  models: Model[];
+  token_usage: TokenUsageSettings;
+}