diff --git a/src/cli/tui/__tests__/token-usage.test.ts b/src/cli/tui/__tests__/token-usage.test.ts new file mode 100644 index 0000000..f2ad1fd --- /dev/null +++ b/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, + }); + }); +}); diff --git a/src/cli/tui/components/footer.tsx b/src/cli/tui/components/footer.tsx index f662ea6..7fbc5f8 100644 --- a/src/cli/tui/components/footer.tsx +++ b/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 ( {agent.model.name} - {formatTokenCount(tokenCount)} tokens + + last input {formatTokenCount(tokenUsage.latestInputTokens)} ยท session{" "} + {formatTokenCount(tokenUsage.sessionTotalTokens)} + ); diff --git a/src/cli/tui/hooks/use-agent-loop.ts b/src/cli/tui/hooks/use-agent-loop.ts index 28dcddd..d43b507 100644 --- a/src/cli/tui/hooks/use-agent-loop.ts +++ b/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; abort: () => void; - tokenCount: number; + tokenUsage: TokenUsageSummary; }; const AgentLoopContext = createContext(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(); } diff --git a/src/cli/tui/token-usage.ts b/src/cli/tui/token-usage.ts new file mode 100644 index 0000000..b8d185a --- /dev/null +++ b/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( + (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"; +}