Clarify token usage in TUI footer

src/cli/tui/__tests__/token-usage.test.ts

@@ -0,0 +1,59 @@
+import { describe, expect, test } from "bun:test";
+
+import type { NonSystemMessage } from "@/foundation";
+
+import { calculateTokenUsage } from "../token-usage";
+
+describe("calculateTokenUsage", () => {
+  test("returns zeroes when no assistant usage exists", () => {
+    const messages: NonSystemMessage[] = [
+      { role: "user", content: [{ type: "text", text: "hello" }] },
+      { role: "tool", content: [{ type: "tool_result", tool_use_id: "tool-1", content: "done" }] },
+    ];
+
+    expect(calculateTokenUsage(messages)).toEqual({
+      latestInputTokens: 0,
+      sessionTotalTokens: 0,
+    });
+  });
+
+  test("uses the latest assistant prompt tokens and cumulative total tokens", () => {
+    const messages: NonSystemMessage[] = [
+      {
+        role: "assistant",
+        content: [{ type: "text", text: "first" }],
+        usage: { promptTokens: 100, completionTokens: 20, totalTokens: 120 },
+      },
+      { role: "user", content: [{ type: "text", text: "next" }] },
+      {
+        role: "assistant",
+        content: [{ type: "text", text: "second" }],
+        usage: { promptTokens: 250, completionTokens: 30, totalTokens: 280 },
+      },
+    ];
+
+    expect(calculateTokenUsage(messages)).toEqual({
+      latestInputTokens: 250,
+      sessionTotalTokens: 400,
+    });
+  });
+
+  test("ignores assistant messages without usage", () => {
+    const messages: NonSystemMessage[] = [
+      {
+        role: "assistant",
+        content: [{ type: "text", text: "first" }],
+        usage: { promptTokens: 40, completionTokens: 10, totalTokens: 50 },
+      },
+      {
+        role: "assistant",
+        content: [{ type: "text", text: "missing usage" }],
+      },
+    ];
+
+    expect(calculateTokenUsage(messages)).toEqual({
+      latestInputTokens: 40,
+      sessionTotalTokens: 50,
+    });
+  });
+});

src/cli/tui/components/footer.tsx

@@ -11,14 +11,17 @@ function formatTokenCount(count: number): string {
 }
 
 export function Footer() {
-  const { agent, tokenCount } = useAgentLoop();
+  const { agent, tokenUsage } = useAgentLoop();
   return (
     <Box paddingX={2} width="100%">
       <Box flexGrow={1} justifyContent="flex-start">
         <Text color={currentTheme.colors.dimText}>{agent.model.name}</Text>
       </Box>
       <Box justifyContent="flex-end">
-        <Text color={currentTheme.colors.dimText}>{formatTokenCount(tokenCount)} tokens</Text>
+        <Text color={currentTheme.colors.dimText}>
+          last input {formatTokenCount(tokenUsage.latestInputTokens)} · session{" "}
+          {formatTokenCount(tokenUsage.sessionTotalTokens)}
+        </Text>
       </Box>
     </Box>
   );

src/cli/tui/hooks/use-agent-loop.ts

@@ -6,6 +6,7 @@ import type { AssistantMessage, NonSystemMessage, UserMessage } from "@/foundati
 
 import type { PromptSubmission, SlashCommand } from "../command-registry";
 import { formatHelp, resolveBuiltinCommand } from "../command-registry";
+import { calculateTokenUsage, type TokenUsageSummary } from "../token-usage";
 
 type AgentLoopState = {
   agent: Agent;
@@ -14,7 +15,7 @@ type AgentLoopState = {
   // eslint-disable-next-line no-unused-vars
   onSubmit: (submission: PromptSubmission) => Promise<void>;
   abort: () => void;
-  tokenCount: number;
+  tokenUsage: TokenUsageSummary;
 };
 
 const AgentLoopContext = createContext<AgentLoopState | null>(null);
@@ -75,8 +76,8 @@ export function AgentLoopProvider({
     agent.abort();
   }, [agent]);
 
-  const tokenCount = useMemo(() => {
-    return calculateTotalTokens(messages);
+  const tokenUsage = useMemo(() => {
+    return calculateTokenUsage(messages);
   }, [messages]);
 
   const onSubmit = useCallback(
@@ -150,9 +151,9 @@ export function AgentLoopProvider({
       messages,
       onSubmit,
       abort,
-      tokenCount,
+      tokenUsage,
     }),
-    [abort, agent, messages, onSubmit, streaming, tokenCount],
+    [abort, agent, messages, onSubmit, streaming, tokenUsage],
   );
 
   return createElement(AgentLoopContext.Provider, { value }, children);
@@ -166,17 +167,6 @@ function useAgentLoopState(): AgentLoopState {
   return state;
 }
 
-function calculateTotalTokens(messages: NonSystemMessage[]): number {
-  return messages.reduce((total, message) => {
-    if (!isAssistantMessage(message)) return total;
-    return total + (message.usage?.totalTokens ?? 0);
-  }, 0);
-}
-
-function isAssistantMessage(message: NonSystemMessage): message is AssistantMessage {
-  return message.role === "assistant";
-}
-
 export function useAgentLoop() {
   return useAgentLoopState();
 }

src/cli/tui/token-usage.ts

@@ -0,0 +1,26 @@
+import type { AssistantMessage, NonSystemMessage } from "@/foundation";
+
+export interface TokenUsageSummary {
+  latestInputTokens: number;
+  sessionTotalTokens: number;
+}
+
+export function calculateTokenUsage(messages: NonSystemMessage[]): TokenUsageSummary {
+  return messages.reduce<TokenUsageSummary>(
+    (summary, message) => {
+      if (!isAssistantMessage(message) || !message.usage) {
+        return summary;
+      }
+
+      return {
+        latestInputTokens: message.usage.promptTokens,
+        sessionTotalTokens: summary.sessionTotalTokens + message.usage.totalTokens,
+      };
+    },
+    { latestInputTokens: 0, sessionTotalTokens: 0 },
+  );
+}
+
+function isAssistantMessage(message: NonSystemMessage): message is AssistantMessage {
+  return message.role === "assistant";
+}