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>

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

@@ -1,29 +1,27 @@
 import type { Message } from "@langchain/langgraph-sdk";
 import { CoinsIcon } from "lucide-react";
 
+import { Badge } from "@/components/ui/badge";
 import { useI18n } from "@/core/i18n/hooks";
-import { formatTokenCount, getUsageMetadata } from "@/core/messages/usage";
+import { accumulateUsage, formatTokenCount } from "@/core/messages/usage";
+import type { TokenDebugStep } from "@/core/messages/usage-model";
 import { cn } from "@/lib/utils";
 
-export function MessageTokenUsage({
+function TokenUsageSummary({
   className,
-  enabled = false,
-  isLoading = false,
-  message,
+  inputTokens,
+  outputTokens,
+  totalTokens,
+  unavailable = false,
 }: {
   className?: string;
-  enabled?: boolean;
-  isLoading?: boolean;
-  message: Message;
+  inputTokens?: number;
+  outputTokens?: number;
+  totalTokens?: number;
+  unavailable?: boolean;
 }) {
   const { t } = useI18n();
 
-  if (!enabled || isLoading || message.type !== "ai") {
-    return null;
-  }
-
-  const usage = getUsageMetadata(message);
-
   return (
     <div
       className={cn(
@@ -35,16 +33,16 @@ export function MessageTokenUsage({
         <CoinsIcon className="size-3" />
         {t.tokenUsage.label}
       </span>
-      {usage ? (
+      {!unavailable ? (
         <>
           <span>
-            {t.tokenUsage.input}: {formatTokenCount(usage.inputTokens)}
+            {t.tokenUsage.input}: {formatTokenCount(inputTokens ?? 0)}
           </span>
           <span>
-            {t.tokenUsage.output}: {formatTokenCount(usage.outputTokens)}
+            {t.tokenUsage.output}: {formatTokenCount(outputTokens ?? 0)}
           </span>
           <span className="font-medium">
-            {t.tokenUsage.total}: {formatTokenCount(usage.totalTokens)}
+            {t.tokenUsage.total}: {formatTokenCount(totalTokens ?? 0)}
           </span>
         </>
       ) : (
@@ -75,17 +73,93 @@ export function MessageTokenUsageList({
     return null;
   }
 
+  const usage = accumulateUsage(aiMessages);
+
   return (
-    <>
-      {aiMessages.map((message, index) => (
-        <MessageTokenUsage
-          className={className}
-          enabled={enabled}
-          isLoading={isLoading}
-          key={message.id ?? index}
-          message={message}
-        />
-      ))}
-    </>
+    <TokenUsageSummary
+      className={className}
+      inputTokens={usage?.inputTokens}
+      outputTokens={usage?.outputTokens}
+      totalTokens={usage?.totalTokens}
+      unavailable={!usage}
+    />
+  );
+}
+
+export function MessageTokenUsageDebugList({
+  className,
+  enabled = false,
+  isLoading = false,
+  steps,
+}: {
+  className?: string;
+  enabled?: boolean;
+  isLoading?: boolean;
+  steps: TokenDebugStep[];
+}) {
+  const { t } = useI18n();
+
+  if (!enabled || isLoading) {
+    return null;
+  }
+
+  if (steps.length === 0) {
+    return null;
+  }
+
+  return (
+    <div className={cn("border-border/60 mt-1 border-t pt-2", className)}>
+      <div className="space-y-2">
+        {steps.map((step) => (
+          <div
+            key={step.id}
+            className="bg-muted/30 border-border/50 flex items-start justify-between gap-3 rounded-md border px-3 py-2"
+          >
+            <div className="min-w-0 flex-1 space-y-1">
+              <div className="text-foreground flex items-center gap-2 text-xs font-medium">
+                <CoinsIcon className="text-muted-foreground size-3" />
+                <span className="truncate">{step.label}</span>
+              </div>
+              {step.secondaryLabels.length > 0 && (
+                <div className="flex flex-wrap gap-1.5">
+                  {step.secondaryLabels.map((label, index) => (
+                    <Badge
+                      key={`${step.id}-${index}-${label}`}
+                      className="px-1.5 py-0 text-[10px] font-normal"
+                      variant="secondary"
+                    >
+                      {label}
+                    </Badge>
+                  ))}
+                </div>
+              )}
+              {step.sharedAttribution && (
+                <div className="text-muted-foreground text-[11px]">
+                  {t.tokenUsage.sharedAttribution}
+                </div>
+              )}
+              <div className="text-muted-foreground text-[11px]">
+                {step.usage ? (
+                  <>
+                    {t.tokenUsage.input}:{" "}
+                    {formatTokenCount(step.usage.inputTokens)}
+                    {" · "}
+                    {t.tokenUsage.output}:{" "}
+                    {formatTokenCount(step.usage.outputTokens)}
+                  </>
+                ) : (
+                  t.tokenUsage.unavailableShort
+                )}
+              </div>
+            </div>
+            <Badge className="shrink-0 font-mono" variant="outline">
+              {step.usage
+                ? `${formatTokenCount(step.usage.totalTokens)} ${t.tokenUsage.label}`
+                : t.tokenUsage.unavailableShort}
+            </Badge>
+          </div>
+        ))}
+      </div>
+    </div>
   );
 }