getcompanion-ai/co-mono
1
0
Fork 0
mirror of https://github.com/getcompanion-ai/co-mono.git synced 2026-04-21 09:01:20 +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 {

28
packages/tui/src/components/text-editor.ts
View file

@ -649,6 +649,8 @@ export class TextEditor implements Component {
// Check if we're in a slash command context
if (beforeCursor.trimStart().startsWith("/")) {
this.handleSlashCommandCompletion();
} else {
this.forceFileAutocomplete();
}
}
@ -656,8 +658,34 @@ export class TextEditor implements Component {
// For now, fall back to regular autocomplete (slash commands)
// This can be extended later to handle command-specific argument completion
this.tryTriggerAutocomplete(true);
}
private forceFileAutocomplete(): void {
if (!this.autocompleteProvider) return;
// Check if provider has the force method
const provider = this.autocompleteProvider as any;
if (!provider.getForceFileSuggestions) {
this.tryTriggerAutocomplete(true);
return;
}
const suggestions = provider.getForceFileSuggestions(
this.state.lines,
this.state.cursorLine,
this.state.cursorCol,
);
if (suggestions && suggestions.items.length > 0) {
this.autocompletePrefix = suggestions.prefix;
this.autocompleteList = new SelectList(suggestions.items, 5);
this.isAutocompleting = true;
} else {
this.cancelAutocomplete();
}
}
private cancelAutocomplete(): void {
this.isAutocompleting = false;
this.autocompleteList = undefined as any;

269
packages/tui/src/tui.ts
View file

@ -272,9 +272,6 @@ export class TUI extends Container {
this.renderInitial(currentRenderCommands);
this.isFirstRender = false;
} else {
// this.executeDifferentialRender(currentRenderCommands, termHeight);
// this.renderDifferential(currentRenderCommands, termHeight);
// this.renderDifferentialSurgical(currentRenderCommands, termHeight);
this.renderLineBased(currentRenderCommands, termHeight);
}
@ -464,272 +461,6 @@ export class TUI extends Container {
this.totalLinesRedrawn += linesRedrawn;
}
private renderDifferentialSurgical(currentCommands: RenderCommand[], termHeight: number): void {
const viewportHeight = termHeight - 1; // Leave one line for cursor
// Build the new lines array
const newLines: string[] = [];
for (const command of currentCommands) {
newLines.push(...command.lines);
}
const totalNewLines = newLines.length;
const totalOldLines = this.previousLines.length;
// Phase 1: Analyze - categorize all changes
let firstChangeOffset = -1;
let hasLineCountChange = false;
let hasStructuralChange = false;
const changedLines: Array<{ lineIndex: number; newContent: string }> = [];
let currentLineOffset = 0;
for (let i = 0; i < Math.max(currentCommands.length, this.previousRenderCommands.length); i++) {
const current = i < currentCommands.length ? currentCommands[i] : null;
const previous = i < this.previousRenderCommands.length ? this.previousRenderCommands[i] : null;
// Structural change: component added/removed/reordered
if (!current || !previous || current.id !== previous.id) {
hasStructuralChange = true;
if (firstChangeOffset === -1) {
firstChangeOffset = currentLineOffset;
}
break;
}
// Line count change
if (current.changed && current.lines.length !== previous.lines.length) {
hasLineCountChange = true;
if (firstChangeOffset === -1) {
firstChangeOffset = currentLineOffset;
}
break;
}
// Content change with same line count - track individual line changes
if (current.changed) {
for (let j = 0; j < current.lines.length; j++) {
const oldLine =
currentLineOffset + j < this.previousLines.length ? this.previousLines[currentLineOffset + j] : "";
const newLine = current.lines[j];
if (oldLine !== newLine) {
changedLines.push({
lineIndex: currentLineOffset + j,
newContent: newLine,
});
if (firstChangeOffset === -1) {
firstChangeOffset = currentLineOffset + j;
}
}
}
}
currentLineOffset += current ? current.lines.length : 0;
}
// If nothing changed, do nothing
if (firstChangeOffset === -1) {
this.previousLines = newLines;
return;
}
// Phase 2: Decision - pick rendering strategy
const contentStartInViewport = Math.max(0, totalOldLines - viewportHeight);
const changePositionInViewport = firstChangeOffset - contentStartInViewport;
let output = "";
let linesRedrawn = 0;
if (changePositionInViewport < 0) {
// Strategy: FULL - change is above viewport, must clear scrollback and re-render all
output = "\x1b[3J\x1b[H"; // Clear scrollback and screen, then home cursor
for (let i = 0; i < newLines.length; i++) {
if (i > 0) output += "\r\n";
output += newLines[i];
}
if (newLines.length > 0) output += "\r\n";
linesRedrawn = newLines.length;
} else if (hasStructuralChange || hasLineCountChange) {
// Strategy: PARTIAL - changes in viewport but with shifts, clear from change to end
// After rendering with a final newline, cursor is one line below the last content line
// So if we have N lines (0 to N-1), cursor is at line N
// To move to line firstChangeOffset, we need to move up (N - firstChangeOffset) lines
// But since cursor is at N (not N-1), we actually need to move up (N - firstChangeOffset) lines
// which is totalOldLines - firstChangeOffset
const cursorLine = totalOldLines; // Cursor is one past the last line
const targetLine = firstChangeOffset;
const linesToMoveUp = cursorLine - targetLine;
if (linesToMoveUp > 0) {
output += `\x1b[${linesToMoveUp}A`;
}
// Clear from cursor to end of screen
// First ensure we're at the beginning of the line
output += "\r";
output += "\x1b[0J"; // Clear from cursor to end of screen
const linesToRender = newLines.slice(firstChangeOffset);
for (let i = 0; i < linesToRender.length; i++) {
if (i > 0) output += "\r\n";
output += linesToRender[i];
}
if (linesToRender.length > 0) output += "\r\n";
linesRedrawn = linesToRender.length;
} else {
// Strategy: SURGICAL - only content changes with same line counts, update only changed lines
// The cursor starts at the line after our last content
let currentCursorLine = totalOldLines;
for (const change of changedLines) {
// Move cursor to the line that needs updating
const linesToMove = currentCursorLine - change.lineIndex;
if (linesToMove > 0) {
output += `\x1b[${linesToMove}A`; // Move up
} else if (linesToMove < 0) {
output += `\x1b[${-linesToMove}B`; // Move down
}
// Clear the line and write new content
output += "\x1b[2K"; // Clear entire line
output += "\r"; // Move to start of line
output += change.newContent;
// Cursor is now at the end of the content on this line
currentCursorLine = change.lineIndex;
linesRedrawn++;
}
// Return cursor to end position
// We need to be on the line after our last content line
// First ensure we're at start of current line
output += "\r";
// Move to last content line
const lastContentLine = totalNewLines - 1;
const linesToMove = lastContentLine - currentCursorLine;
if (linesToMove > 0) {
output += `\x1b[${linesToMove}B`;
} else if (linesToMove < 0) {
output += `\x1b[${-linesToMove}A`;
}
// Now add final newline to position cursor on next line
output += "\r\n";
}
this.terminal.write(output);
// Save what we rendered
this.previousLines = newLines;
this.totalLinesRedrawn += linesRedrawn;
}
// biome-ignore lint/correctness/noUnusedPrivateClassMembers: Keeping this around as reference for LLM
private renderDifferential(currentCommands: RenderCommand[], termHeight: number): void {
const viewportHeight = termHeight - 1; // Leave one line for cursor
// Build the new lines array
const newLines: string[] = [];
for (const command of currentCommands) {
newLines.push(...command.lines);
}
const totalNewLines = newLines.length;
const totalOldLines = this.previousLines.length;
// Find the first line that changed
let firstChangedLineOffset = -1;
let currentLineOffset = 0;
for (let i = 0; i < currentCommands.length; i++) {
const current = currentCommands[i];
const previous = i < this.previousRenderCommands.length ? this.previousRenderCommands[i] : null;
// Check if this is a new component or component was removed/reordered
if (!previous || previous.id !== current.id) {
firstChangedLineOffset = currentLineOffset;
break;
}
// Check if component content or size changed
if (current.changed) {
firstChangedLineOffset = currentLineOffset;
break;
}
currentLineOffset += current.lines.length;
}
// Also check if we have fewer components now (components removed from end)
if (firstChangedLineOffset === -1 && currentCommands.length < this.previousRenderCommands.length) {
firstChangedLineOffset = currentLineOffset;
}
// If nothing changed, do nothing
if (firstChangedLineOffset === -1) {
this.previousLines = newLines;
return;
}
// Calculate where the first change is relative to the viewport
// If our content exceeds viewport, some is in scrollback
const contentStartInViewport = Math.max(0, totalOldLines - viewportHeight);
const changePositionInViewport = firstChangedLineOffset - contentStartInViewport;
let output = "";
let linesRedrawn = 0;
if (changePositionInViewport < 0) {
// The change is above the viewport - we cannot reach it with cursor
// MUST do full re-render
output = "\x1b[3J\x1b[H"; // Clear scrollback and screen, then home cursor
// Render ALL lines
for (let i = 0; i < newLines.length; i++) {
if (i > 0) output += "\r\n";
output += newLines[i];
}
// Add final newline
if (newLines.length > 0) output += "\r\n";
linesRedrawn = newLines.length;
} else {
// The change is in the viewport - we can update from there
// Calculate how many lines up to move from current cursor position
const linesToMoveUp = totalOldLines - firstChangedLineOffset;
if (linesToMoveUp > 0) {
output += `\x1b[${linesToMoveUp}A`;
}
// Clear from here to end of screen
output += "\x1b[0J";
// Render everything from the first change onwards
const linesToRender = newLines.slice(firstChangedLineOffset);
for (let i = 0; i < linesToRender.length; i++) {
if (i > 0) output += "\r\n";
output += linesToRender[i];
}
// Add final newline
if (linesToRender.length > 0) output += "\r\n";
linesRedrawn = linesToRender.length;
}
this.terminal.write(output);
// Save what we rendered
this.previousLines = newLines;
this.totalLinesRedrawn += linesRedrawn;
}
private handleResize(): void {
// Clear screen and reset
this.terminal.write("\x1b[2J\x1b[H\x1b[?25l");