fix(ui): correctly display token usage across sessions (kagent-dev#1509)

When testing kagent-dev#1500, I struggled to visually demonstrate the affect of
the duplicated session history (aka context bloat) when chatting with
agents. The current implementation of "Usage" within the UI (introduced
in kagent-dev#679) uses `Math.max` independently per field (`total`, `prompt`,
`completion`) across all tasks in a session. This produces incoherent
values where `total` does not necessarily match `prompt` + `completion`,
since each field[^1] (most notably `prompt`) can peak during differing
tasks. And, it doesn't correctly update to reflect usage across multiple
API calls in a multi-turn conversation. All we see is which API call has
the greatest of each type of token usage.

This PR addresses those shortcomings by:

- Updating the existing token usage display to reflect session total
across multiple turns
- Show per-call token usage tooltip on tool call cards and text messages

[^1]: not 100% certain here but I think that fields other than input
could also peak at different levels within cached prompt scenarios.

Signed-off-by: Brian Fox <878612+onematchfox@users.noreply.github.com>

Author: onematchfox

python/packages/kagent-adk/src/kagent/adk/_agent_executor.py

@@ -53,7 +53,7 @@
 )
 
 from ._mcp_toolset import is_anyio_cross_task_cancel_scope_error
-from .converters.event_converter import convert_event_to_a2a_events
+from .converters.event_converter import convert_event_to_a2a_events, serialize_metadata_value
 from .converters.part_converter import convert_a2a_part_to_genai_part, convert_genai_part_to_a2a_part
 from .converters.request_converter import convert_a2a_request_to_adk_run_args
 
@@ -585,6 +585,7 @@ async def _handle_request(
         # For streaming A2A update events, the invocation_id is added through event converter
         # This adds the invocation_id of the run to the metadata of the FINAL event (completed or failed)
         real_invocation_id: str | None = None
+        last_usage_metadata = None
 
         task_result_aggregator = TaskResultAggregator()
         async with Aclosing(runner.run_async(**run_args)) as agen:
@@ -595,6 +596,12 @@ async def _handle_request(
                     real_invocation_id = event_inv_id
                     run_metadata[get_kagent_metadata_key("invocation_id")] = real_invocation_id
 
+                # Track the last usage_metadata so it can be included in the final
+                # event's run_metadata. The A2A task_manager merges run_metadata into
+                # task.metadata, making it available to callers (e.g. KAgentRemoteA2ATool).
+                if getattr(adk_event, "usage_metadata", None) is not None:
+                    last_usage_metadata = adk_event.usage_metadata
+
                 for a2a_event in convert_event_to_a2a_events(
                     adk_event, invocation_context, context.task_id, context.context_id
                 ):
@@ -608,6 +615,11 @@ async def _handle_request(
                 if getattr(adk_event, "long_running_tool_ids", None):
                     break
 
+        # Attach the last LLM usage to run_metadata so the A2A task_manager
+        # merges it into task.metadata on the completed Task object.
+        if last_usage_metadata is not None:
+            run_metadata[get_kagent_metadata_key("usage_metadata")] = serialize_metadata_value(last_usage_metadata)
+
         # publish the task result event - this is final
         if (
             task_result_aggregator.task_state == TaskState.working

python/packages/kagent-adk/src/kagent/adk/_remote_a2a_tool.py

@@ -83,6 +83,21 @@ def _extract_text_from_task(task: Task) -> str:
     return ""
 
 
+def _extract_usage_from_task(task: Task) -> Optional[dict]:
+    """Extract kagent_usage_metadata from a completed task.
+
+    The A2A task_manager merges the final TaskStatusUpdateEvent.metadata into
+    task.metadata. The agent executor now adds the last LLM invocation's
+    usage_metadata to run_metadata before publishing the final event, so it
+    is available here for non-streaming callers like KAgentRemoteA2ATool.
+    """
+    if task.metadata:
+        usage = task.metadata.get("kagent_usage_metadata")
+        if usage and isinstance(usage, dict):
+            return usage
+    return None
+
+
 class KAgentRemoteA2ATool(BaseTool):
     """A tool that calls a remote A2A agent and propagates HITL state."""
 
@@ -215,8 +230,13 @@ async def _handle_first_call(self, args: dict[str, Any], tool_context: ToolConte
             error_text = _extract_text_from_task(task)
             return error_text or f"Remote agent '{self.name}' failed."
 
-        # completed or any other terminal state
-        return _extract_text_from_task(task) or ""
+        # completed — include the sub-agent's final LLM usage from task.metadata
+        # so the parent can display it on the AgentCall card in the UI.
+        result_text = _extract_text_from_task(task)
+        usage = _extract_usage_from_task(task)
+        if usage:
+            return {"result": result_text, "kagent_usage_metadata": usage}
+        return result_text or ""
 
     def _handle_input_required(self, task: Task, tool_context: ToolContext) -> dict[str, Any]:
         """Handle a subagent that returned input_required (HITL).
@@ -351,7 +371,11 @@ async def _handle_resume(self, tool_context: ToolContext) -> Any:
             error_text = _extract_text_from_task(task)
             return error_text or f"Remote agent '{subagent_name}' failed after resume."
 
-        return _extract_text_from_task(task) or ""
+        result_text = _extract_text_from_task(task)
+        usage = _extract_usage_from_task(task)
+        if usage:
+            return {"result": result_text, "kagent_usage_metadata": usage}
+        return result_text or ""
 
     @staticmethod
     def _extract_text_from_message(message: A2AMessage) -> str:

python/packages/kagent-adk/src/kagent/adk/converters/event_converter.py

@@ -33,7 +33,7 @@
 logger = logging.getLogger("kagent_adk." + __name__)
 
 
-def _serialize_metadata_value(value: Any) -> str:
+def serialize_metadata_value(value: Any) -> str:
     """Safely serializes metadata values to string format.
 
     Args:
@@ -90,7 +90,7 @@ def _get_context_metadata(event: Event, invocation_context: InvocationContext) -
 
         for field_name, field_value in optional_fields:
             if field_value is not None:
-                metadata[get_kagent_metadata_key(field_name)] = _serialize_metadata_value(field_value)
+                metadata[get_kagent_metadata_key(field_name)] = serialize_metadata_value(field_value)
 
         return metadata
 

ui/src/components/ToolDisplay.tsx

@@ -1,10 +1,11 @@
 import { useState } from "react";
-import { FunctionCall } from "@/types";
+import { FunctionCall, TokenStats } from "@/types";
 import { ScrollArea } from "@radix-ui/react-scroll-area";
 import { FunctionSquare, CheckCircle, Clock, Code, ChevronUp, ChevronDown, Loader2, Text, Check, Copy, AlertCircle, ShieldAlert } from "lucide-react";
 import { Button } from "@/components/ui/button";
 import { Textarea } from "@/components/ui/textarea";
 import { Card, CardHeader, CardTitle, CardContent } from "@/components/ui/card";
+import TokenStatsTooltip from "@/components/chat/TokenStatsTooltip";
 import { convertToUserFriendlyName } from "@/lib/utils";
 
 export type ToolCallStatus = "requested" | "executing" | "completed" | "pending_approval" | "approved" | "rejected";
@@ -23,9 +24,10 @@ interface ToolDisplayProps {
   subagentName?: string;
   onApprove?: () => void;
   onReject?: (reason?: string) => void;
+  tokenStats?: TokenStats;
 }
 
-const ToolDisplay = ({ call, result, status = "requested", isError = false, isDecided = false, subagentName, onApprove, onReject }: ToolDisplayProps) => {
+const ToolDisplay = ({ call, result, status = "requested", isError = false, isDecided = false, subagentName, onApprove, onReject, tokenStats }: ToolDisplayProps) => {
   const [areArgumentsExpanded, setAreArgumentsExpanded] = useState(status === "pending_approval");
   const [areResultsExpanded, setAreResultsExpanded] = useState(false);
   const [isCopied, setIsCopied] = useState(false);
@@ -166,7 +168,8 @@ const ToolDisplay = ({ call, result, status = "requested", isError = false, isDe
           )}
           <div className="font-light">{call.id}</div>
         </CardTitle>
-        <div className="flex justify-center items-center text-xs">
+        <div className="flex items-center gap-2 text-xs">
+          {tokenStats && <TokenStatsTooltip stats={tokenStats} />}
           {getStatusDisplay()}
         </div>
       </CardHeader>

ui/src/components/chat/AgentCallDisplay.tsx

@@ -1,9 +1,10 @@
 import { useMemo, useState } from "react";
-import { FunctionCall } from "@/types";
+import { FunctionCall, TokenStats } from "@/types";
 import { Card, CardHeader, CardTitle, CardContent } from "@/components/ui/card";
 import { convertToUserFriendlyName } from "@/lib/utils";
 import { ChevronDown, ChevronUp, MessageSquare, Loader2, AlertCircle, CheckCircle } from "lucide-react";
 import KagentLogo from "../kagent-logo";
+import TokenStatsTooltip from "@/components/chat/TokenStatsTooltip";
 
 export type AgentCallStatus = "requested" | "executing" | "completed";
 
@@ -15,9 +16,10 @@ interface AgentCallDisplayProps {
   };
   status?: AgentCallStatus;
   isError?: boolean;
+  tokenStats?: TokenStats;
 }
 
-const AgentCallDisplay = ({ call, result, status = "requested", isError = false }: AgentCallDisplayProps) => {
+const AgentCallDisplay = ({ call, result, status = "requested", isError = false, tokenStats }: AgentCallDisplayProps) => {
   const [areInputsExpanded, setAreInputsExpanded] = useState(false);
   const [areResultsExpanded, setAreResultsExpanded] = useState(false);
 
@@ -78,7 +80,8 @@ const AgentCallDisplay = ({ call, result, status = "requested", isError = false
           </div>
           <div className="font-light">{call.id}</div>
         </CardTitle>
-        <div className="flex justify-center items-center text-xs">
+        <div className="flex items-center gap-2 text-xs">
+          {tokenStats && <TokenStatsTooltip stats={tokenStats} />}
           {getStatusDisplay()}
         </div>
       </CardHeader>

ui/src/components/chat/ChatInterface.tsx

@@ -1,7 +1,7 @@
 "use client";
 
 import type React from "react";
-import { useState, useRef, useEffect } from "react";
+import { useState, useRef, useEffect, useMemo } from "react";
 import { ArrowBigUp, X, Loader2, Mic, Square } from "lucide-react";
 import { Button } from "@/components/ui/button";
 import {
@@ -15,7 +15,7 @@ import { Textarea } from "@/components/ui/textarea";
 import { ScrollArea } from "@/components/ui/scroll-area";
 import ChatMessage from "@/components/chat/ChatMessage";
 import StreamingMessage from "./StreamingMessage";
-import TokenStatsDisplay from "./TokenStats";
+import SessionTokenStatsDisplay from "@/components/chat/TokenStats";
 import type { TokenStats, Session, ChatStatus, ToolDecision } from "@/types";
 import StatusDisplay from "./StatusDisplay";
 import { createSession, getSessionTasks, checkSessionExists } from "@/app/actions/sessions";
@@ -39,11 +39,6 @@ export default function ChatInterface({ selectedAgentName, selectedNamespace, se
   const router = useRouter();
   const containerRef = useRef<HTMLDivElement>(null);
   const [currentInputMessage, setCurrentInputMessage] = useState("");
-  const [tokenStats, setTokenStats] = useState<TokenStats>({
-    total: 0,
-    input: 0,
-    output: 0,
-  });
 
   const [chatStatus, setChatStatus] = useState<ChatStatus>("ready");
 
@@ -58,6 +53,9 @@ export default function ChatInterface({ selectedAgentName, selectedNamespace, se
   const [sessionNotFound, setSessionNotFound] = useState<boolean>(false);
   const isCreatingSessionRef = useRef<boolean>(false);
   const [isFirstMessage, setIsFirstMessage] = useState<boolean>(!sessionId);
+  const [sessionStats, setSessionStats] = useState<TokenStats>({ total: 0, prompt: 0, completion: 0 });
+  // Mutable ref so pendingTurnStats survives re-renders between A2A stream events
+  const pendingTurnStatsRef = useRef<TokenStats | undefined>(undefined);
   const [pendingDecisions, setPendingDecisions] = useState<Record<string, ToolDecision>>({});
   const pendingDecisionsRef = useRef<Record<string, ToolDecision>>({});
   /** Per-tool rejection reasons collected as the user rejects individual tools. */
@@ -78,25 +76,27 @@ export default function ChatInterface({ selectedAgentName, selectedNamespace, se
     },
   });
 
-  const { handleMessageEvent } = createMessageHandlers({
+  const { handleMessageEvent } = useMemo(() => createMessageHandlers({
     setMessages: setStreamingMessages,
     setIsStreaming,
     setStreamingContent,
-    setTokenStats,
     setChatStatus,
+    setSessionStats,
+    pendingTurnStats: pendingTurnStatsRef,
     agentContext: {
       namespace: selectedNamespace,
       agentName: selectedAgentName
     }
-  });
+  }), [selectedNamespace, selectedAgentName]);
 
   useEffect(() => {
     async function initializeChat() {
-      setTokenStats({ total: 0, input: 0, output: 0 });
+      setSessionStats({ total: 0, prompt: 0, completion: 0 });
       setStreamingMessages([]);
       setPendingDecisions({});
       pendingDecisionsRef.current = {};
       pendingRejectionReasonsRef.current = {};
+      pendingTurnStatsRef.current = undefined;
 
       // Skip completely if this is a first message session creation flow
       if (isFirstMessage || isCreatingSessionRef.current) {
@@ -129,11 +129,11 @@ export default function ChatInterface({ selectedAgentName, selectedNamespace, se
         }
         if (!messagesResponse.data || messagesResponse?.data?.length === 0) {
           setStoredMessages([]);
-          setTokenStats({ total: 0, input: 0, output: 0 });
+          setSessionStats({ total: 0, prompt: 0, completion: 0 });
         }
         else {
           const extractedMessages = extractMessagesFromTasks(messagesResponse.data);
-          const extractedTokenStats = extractTokenStatsFromTasks(messagesResponse.data);
+          setSessionStats(extractTokenStatsFromTasks(messagesResponse.data));
 
           // Resolved approvals are already inline in extractedMessages (with
           // approved/rejected badges). Only pending approvals need appending.
@@ -144,7 +144,6 @@ export default function ChatInterface({ selectedAgentName, selectedNamespace, se
               ? [...extractedMessages, ...pendingApprovalMessages]
               : extractedMessages
           );
-          setTokenStats(extractedTokenStats);
 
           if (hasPendingApproval) {
             setChatStatus("input_required");
@@ -192,6 +191,7 @@ export default function ChatInterface({ selectedAgentName, selectedNamespace, se
     setPendingDecisions({});
     pendingDecisionsRef.current = {};
     pendingRejectionReasonsRef.current = {};
+    pendingTurnStatsRef.current = undefined;
 
     // For new sessions or when no stored messages exist, show the user message immediately
     const userMessage: Message = {
@@ -682,7 +682,7 @@ export default function ChatInterface({ selectedAgentName, selectedNamespace, se
       <div className="w-full sticky bg-secondary bottom-0 md:bottom-2 rounded-none md:rounded-lg p-4 border  overflow-hidden transition-all duration-300 ease-in-out">
         <div className="flex items-center justify-between mb-4">
           <StatusDisplay chatStatus={chatStatus} />
-          <TokenStatsDisplay stats={tokenStats} />
+          {sessionStats.total > 0 && <SessionTokenStatsDisplay stats={sessionStats} />}
         </div>
 
         <form onSubmit={handleSendMessage}>

ui/src/components/chat/ChatMessage.tsx

@@ -4,6 +4,8 @@ import ToolCallDisplay from "@/components/chat/ToolCallDisplay";
 import AskUserDisplay, { AskUserQuestion } from "@/components/chat/AskUserDisplay";
 import KagentLogo from "../kagent-logo";
 import { ThumbsUp, ThumbsDown } from "lucide-react";
+import TokenStatsTooltip from "@/components/chat/TokenStatsTooltip";
+import type { TokenStats } from "@/types";
 import { useState } from "react";
 import { FeedbackDialog } from "./FeedbackDialog";
 import { toast } from "sonner";
@@ -28,10 +30,13 @@ export default function ChatMessage({ message, allMessages, agentContext, onAppr
   const [feedbackDialogOpen, setFeedbackDialogOpen] = useState(false);
   const [isPositiveFeedback, setIsPositiveFeedback] = useState(true);
 
+  if (!message) return null;
+
   const textParts = message.parts?.filter(part => part.kind === "text") || [];
   const content = textParts.map(part => (part as TextPart).text).join("");
 
   const source = message.role === "user" ? "user" : "assistant";
+  const tokenStats = (message.metadata as Record<string, unknown> | undefined)?.tokenStats as TokenStats | undefined;
   const messageId = message.messageId;
 
   // Extract agent name from metadata for display
@@ -68,10 +73,6 @@ export default function ChatMessage({ message, allMessages, agentContext, onAppr
     return a & a;
   }, 0)) : 0;
 
-  if (!message) {
-    return null;
-  }
-
   const metadata = message.metadata as ADKMetadata;
   const originalType = metadata?.originalType;
 
@@ -168,22 +169,27 @@ export default function ChatMessage({ message, allMessages, agentContext, onAppr
         <div className="text-xs font-bold">{displayName}</div>
       </div> : <div className="text-xs font-bold">{displayName}</div>}
       <TruncatableText content={String(content)} className="break-all text-primary-foreground" />
-      {source !== "user" && messageId !== undefined && (
+      {source !== "user" && (
         <div className="flex mt-2 justify-end items-center gap-2">
-          <button
-            onClick={() => handleFeedback(true)}
-            className="p-1 rounded-full hover:bg-gray-200 dark:hover:bg-gray-700 transition-colors"
-            aria-label="Thumbs up"
-          >
-            <ThumbsUp className="w-4 h-4" />
-          </button>
-          <button
-            onClick={() => handleFeedback(false)}
-            className="p-1 rounded-full hover:bg-gray-200 dark:hover:bg-gray-700 transition-colors"
-            aria-label="Thumbs down"
-          >
-            <ThumbsDown className="w-4 h-4" />
-          </button>
+          {tokenStats && <TokenStatsTooltip stats={tokenStats} />}
+          {messageId !== undefined && (
+            <>
+              <button
+                onClick={() => handleFeedback(true)}
+                className="p-1 rounded-full hover:bg-gray-200 dark:hover:bg-gray-700 transition-colors"
+                aria-label="Thumbs up"
+              >
+                <ThumbsUp className="w-4 h-4" />
+              </button>
+              <button
+                onClick={() => handleFeedback(false)}
+                className="p-1 rounded-full hover:bg-gray-200 dark:hover:bg-gray-700 transition-colors"
+                aria-label="Thumbs down"
+              >
+                <ThumbsDown className="w-4 h-4" />
+              </button>
+            </>
+          )}
         </div>
       )}
     </div>

ui/src/components/chat/TokenStats.tsx

@@ -1,23 +1,23 @@
 import { ArrowLeft, ArrowRightFromLine } from "lucide-react";
 import { TokenStats } from "@/types";
 
-interface TokenStatsDisplayProps {
+interface SessionTokenStatsDisplayProps {
   stats: TokenStats;
 }
 
-export default function TokenStatsDisplay({ stats }: TokenStatsDisplayProps) {
+export default function SessionTokenStatsDisplay({ stats }: SessionTokenStatsDisplayProps) {
   return (
     <div className="flex items-center gap-2 text-xs">
       <span>Usage: </span>
       <span>{stats.total}</span>
       <div className="flex items-center gap-2">
         <div className="flex items-center gap-1">
           <ArrowLeft className="h-3 w-3" />
-          <span>{stats.input}</span>
+          <span>{stats.prompt}</span>
         </div>
         <div className="flex items-center gap-1">
           <ArrowRightFromLine className="h-3 w-3" />
-          <span>{stats.output}</span>
+          <span>{stats.completion}</span>
         </div>
       </div>
     </div>