getcompanion-ai/co-mono
1
0
Fork 0
mirror of https://github.com/getcompanion-ai/co-mono.git synced 2026-04-17 09:02:08 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

feat(agent): Add /tokens command for cumulative token usage tracking

Browse source 
Added /tokens slash command to TUI that displays session-wide token statistics.
Key changes:
- Fixed SessionManager to accumulate token usage instead of storing only last event
- Added cumulative token tracking to TUI renderer alongside per-request totals
- Implemented slash command infrastructure with /tokens autocomplete support
- Fixed file autocompletion that was missing from Tab key handling
- Clean minimal display format showing input/output/reasoning/cache/tool counts

The /tokens command shows:
Total usage
   input: 1,234
   output: 567
   reasoning: 89
   cache read: 100
   cache write: 50
   tool calls: 2
This commit is contained in:
Mario Zechner 2025-08-11 15:43:48 +02:00
parent 7e3b94ade6
commit e21a46e68f
10 changed files with 303 additions and 283 deletions
Download patch file Download diff file Expand all files Collapse all files

72
packages/agent/src/renderers/tui-renderer.ts
View file

@ -2,6 +2,7 @@ import {
CombinedAutocompleteProvider,
Container,
MarkdownComponent,
SlashCommand,
TextComponent,
TextEditor,
TUI,
@ -63,6 +64,13 @@ export class TuiRenderer implements AgentEventReceiver {
private lastCacheWriteTokens = 0;
private lastReasoningTokens = 0;
private toolCallCount = 0;
// Cumulative token tracking
private cumulativeInputTokens = 0;
private cumulativeOutputTokens = 0;
private cumulativeCacheReadTokens = 0;
private cumulativeCacheWriteTokens = 0;
private cumulativeReasoningTokens = 0;
private cumulativeToolCallCount = 0;
private tokenStatusComponent: TextComponent | null = null;
constructor() {
@ -74,7 +82,12 @@ export class TuiRenderer implements AgentEventReceiver {
// Setup autocomplete for file paths and slash commands
const autocompleteProvider = new CombinedAutocompleteProvider(
[],
[
{
name: "tokens",
description: "Show cumulative token usage for this session",
},
],
process.cwd(), // Base directory for file path completion
);
this.editor.setAutocompleteProvider(autocompleteProvider);
@ -148,6 +161,17 @@ export class TuiRenderer implements AgentEventReceiver {
text = text.trim();
if (!text) return;
// Handle slash commands
if (text.startsWith("/")) {
const [command, ...args] = text.slice(1).split(" ");
if (command === "tokens") {
this.showTokenUsage();
return;
}
// Unknown slash command, ignore
return;
}
if (this.onInputCallback) {
this.onInputCallback(text);
}
@ -192,6 +216,7 @@ export class TuiRenderer implements AgentEventReceiver {
case "tool_call":
this.toolCallCount++;
this.cumulativeToolCallCount++;
this.updateTokenDisplay();
this.chatContainer.addChild(new TextComponent(chalk.yellow(`[tool] ${event.name}(${event.args})`)));
break;
@ -255,6 +280,14 @@ export class TuiRenderer implements AgentEventReceiver {
this.lastCacheReadTokens = event.cacheReadTokens;
this.lastCacheWriteTokens = event.cacheWriteTokens;
this.lastReasoningTokens = event.reasoningTokens;
// Accumulate cumulative totals
this.cumulativeInputTokens += event.inputTokens;
this.cumulativeOutputTokens += event.outputTokens;
this.cumulativeCacheReadTokens += event.cacheReadTokens;
this.cumulativeCacheWriteTokens += event.cacheWriteTokens;
this.cumulativeReasoningTokens += event.reasoningTokens;
this.updateTokenDisplay();
break;
@ -282,21 +315,21 @@ export class TuiRenderer implements AgentEventReceiver {
this.tokenContainer.clear();
// Build token display text
let tokenText = chalk.dim(`↑${this.lastInputTokens.toLocaleString()} ↓${this.lastOutputTokens.toLocaleString()}`);
let tokenText = chalk.dim(`↑ ${this.lastInputTokens.toLocaleString()} ↓ ${this.lastOutputTokens.toLocaleString()}`);
// Add reasoning tokens if present
if (this.lastReasoningTokens > 0) {
tokenText += chalk.dim(` ⚡${this.lastReasoningTokens.toLocaleString()}`);
tokenText += chalk.dim(` ⚡ ${this.lastReasoningTokens.toLocaleString()}`);
}
// Add cache info if available
if (this.lastCacheReadTokens > 0 || this.lastCacheWriteTokens > 0) {
const cacheText: string[] = [];
if (this.lastCacheReadTokens > 0) {
cacheText.push(`⟲${this.lastCacheReadTokens.toLocaleString()}`);
cacheText.push(` cache read: ${this.lastCacheReadTokens.toLocaleString()}`);
}
if (this.lastCacheWriteTokens > 0) {
cacheText.push(`⟳${this.lastCacheWriteTokens.toLocaleString()}`);
cacheText.push(` cache write: ${this.lastCacheWriteTokens.toLocaleString()}`);
}
tokenText += chalk.dim(` (${cacheText.join(" ")})`);
}
@ -346,6 +379,35 @@ export class TuiRenderer implements AgentEventReceiver {
this.ui.requestRender();
}
private showTokenUsage(): void {
let tokenText = chalk.dim(`Total usage\n input: ${this.cumulativeInputTokens.toLocaleString()}\n output: ${this.cumulativeOutputTokens.toLocaleString()}`);
if (this.cumulativeReasoningTokens > 0) {
tokenText += chalk.dim(`\n reasoning: ${this.cumulativeReasoningTokens.toLocaleString()}`);
}
if (this.cumulativeCacheReadTokens > 0 || this.cumulativeCacheWriteTokens > 0) {
const cacheText: string[] = [];
if (this.cumulativeCacheReadTokens > 0) {
cacheText.push(`\n cache read: ${this.cumulativeCacheReadTokens.toLocaleString()}`);
}
if (this.cumulativeCacheWriteTokens > 0) {
cacheText.push(`\n cache right: ${this.cumulativeCacheWriteTokens.toLocaleString()}`);
}
tokenText += chalk.dim(` ${cacheText.join(" ")}`);
}
if (this.cumulativeToolCallCount > 0) {
tokenText += chalk.dim(`\n tool calls: ${this.cumulativeToolCallCount}`);
}
const tokenSummary = new TextComponent(chalk.italic(tokenText), { bottom: 1 });
this.chatContainer.addChild(tokenSummary);
this.ui.requestRender();
}
stop(): void {
if (this.currentLoadingAnimation) {
this.currentLoadingAnimation.stop();

12
packages/agent/src/session-manager.ts
View file

@ -156,7 +156,17 @@ export class SessionManager implements AgentEventReceiver {
const eventEntry: SessionEvent = entry as SessionEvent;
events.push(eventEntry);
if (eventEntry.event.type === "token_usage") {
totalUsage = entry.event as Extract<AgentEvent, { type: "token_usage" }>;
const usage = entry.event as Extract<AgentEvent, { type: "token_usage" }>;
if (!totalUsage) {
totalUsage = { ...usage };
} else {
totalUsage.inputTokens += usage.inputTokens;
totalUsage.outputTokens += usage.outputTokens;
totalUsage.totalTokens += usage.totalTokens;
totalUsage.cacheReadTokens += usage.cacheReadTokens;
totalUsage.cacheWriteTokens += usage.cacheWriteTokens;
totalUsage.reasoningTokens += usage.reasoningTokens;
}
}
}
} catch {