feat: add token consumption tracking per task

- Add TokenUsage interface to track input/output tokens and cost
- Capture usage from Claude SDK result messages in auto-mode-service
- Accumulate token usage across multiple streams and follow-up sessions
- Store token usage in feature.json for persistence
- Display token count and cost on Kanban cards with tooltip breakdown
- Add token usage summary header in log viewer

Closes #166

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
This commit is contained in:
Kacper
2025-12-20 03:05:13 +01:00
parent d104a24446
commit 3559e0104c
7 changed files with 218 additions and 3 deletions

View File

@@ -48,6 +48,29 @@ export interface ContentBlock {
content?: string;
}
/**
* Token usage statistics from SDK execution
*/
export interface TokenUsage {
inputTokens: number;
outputTokens: number;
cacheReadInputTokens: number;
cacheCreationInputTokens: number;
totalTokens: number;
costUSD: number;
}
/**
* Per-model usage breakdown from SDK result
*/
export interface ModelUsageData {
inputTokens: number;
outputTokens: number;
cacheReadInputTokens: number;
cacheCreationInputTokens: number;
costUSD: number;
}
/**
* Message returned by a provider (matches Claude SDK streaming format)
*/
@@ -62,6 +85,15 @@ export interface ProviderMessage {
result?: string;
error?: string;
parent_tool_use_id?: string | null;
// Token usage fields (present in result messages)
usage?: {
input_tokens: number;
output_tokens: number;
cache_read_input_tokens?: number;
cache_creation_input_tokens?: number;
};
total_cost_usd?: number;
modelUsage?: Record<string, ModelUsageData>;
}
/**

View File

@@ -10,7 +10,7 @@
*/
import { ProviderFactory } from "../providers/provider-factory.js";
import type { ExecuteOptions } from "../providers/types.js";
import type { ExecuteOptions, TokenUsage } from "../providers/types.js";
import { exec } from "child_process";
import { promisify } from "util";
import path from "path";
@@ -1878,6 +1878,30 @@ This mock response was generated because AUTOMAKER_MOCK_AGENT=true was set.
: "";
let specDetected = false;
// Token usage accumulator - tracks total usage across all streams
const accumulatedUsage: TokenUsage = {
inputTokens: 0,
outputTokens: 0,
cacheReadInputTokens: 0,
cacheCreationInputTokens: 0,
totalTokens: 0,
costUSD: 0,
};
// Helper to accumulate usage from a result message
const accumulateUsage = (msg: { usage?: { input_tokens: number; output_tokens: number; cache_read_input_tokens?: number; cache_creation_input_tokens?: number }; total_cost_usd?: number }): void => {
if (msg.usage) {
accumulatedUsage.inputTokens += msg.usage.input_tokens || 0;
accumulatedUsage.outputTokens += msg.usage.output_tokens || 0;
accumulatedUsage.cacheReadInputTokens += msg.usage.cache_read_input_tokens || 0;
accumulatedUsage.cacheCreationInputTokens += msg.usage.cache_creation_input_tokens || 0;
accumulatedUsage.totalTokens = accumulatedUsage.inputTokens + accumulatedUsage.outputTokens;
}
if (msg.total_cost_usd !== undefined) {
accumulatedUsage.costUSD = msg.total_cost_usd;
}
};
// Agent output goes to .automaker directory
// Note: We use projectPath here, not workDir, because workDir might be a worktree path
const featureDirForOutput = getFeatureDir(projectPath, featureId);
@@ -2101,6 +2125,7 @@ After generating the revised spec, output:
throw new Error(msg.error || "Error during plan revision");
} else if (msg.type === "result" && msg.subtype === "success") {
revisionText += msg.result || "";
accumulateUsage(msg);
}
}
@@ -2238,6 +2263,7 @@ After generating the revised spec, output:
} else if (msg.type === "result" && msg.subtype === "success") {
taskOutput += msg.result || "";
responseText += msg.result || "";
accumulateUsage(msg);
}
}
@@ -2318,6 +2344,7 @@ Implement all the changes described in the plan above.`;
throw new Error(msg.error || "Unknown error during implementation");
} else if (msg.type === "result" && msg.subtype === "success") {
responseText += msg.result || "";
accumulateUsage(msg);
}
}
}
@@ -2365,6 +2392,8 @@ Implement all the changes described in the plan above.`;
// The msg.result is just a summary which would lose all tool use details
// Just ensure final write happens
scheduleWrite();
// Capture token usage from the result message
accumulateUsage(msg);
}
}
@@ -2374,6 +2403,31 @@ Implement all the changes described in the plan above.`;
}
// Final write - ensure all accumulated content is saved
await writeToFile();
// Save token usage to the feature if any tokens were consumed
if (accumulatedUsage.totalTokens > 0) {
try {
// Load existing feature to check for previous token usage
const existingFeature = await this.featureLoader.get(projectPath, featureId);
if (existingFeature) {
// If feature already has token usage, add to it (for follow-ups)
const existingUsage = existingFeature.tokenUsage;
if (existingUsage) {
accumulatedUsage.inputTokens += existingUsage.inputTokens;
accumulatedUsage.outputTokens += existingUsage.outputTokens;
accumulatedUsage.cacheReadInputTokens += existingUsage.cacheReadInputTokens;
accumulatedUsage.cacheCreationInputTokens += existingUsage.cacheCreationInputTokens;
accumulatedUsage.totalTokens = accumulatedUsage.inputTokens + accumulatedUsage.outputTokens;
accumulatedUsage.costUSD += existingUsage.costUSD;
}
}
await this.featureLoader.update(projectPath, featureId, { tokenUsage: accumulatedUsage });
console.log(`[AutoMode] Token usage:`, accumulatedUsage);
console.log(`[AutoMode] Saved token usage for ${featureId}: ${accumulatedUsage.totalTokens} tokens, $${accumulatedUsage.costUSD.toFixed(4)}`);
} catch (error) {
console.error(`[AutoMode] Failed to save token usage for ${featureId}:`, error);
}
}
}
private async executeFeatureWithContext(

View File

@@ -11,6 +11,7 @@ import {
getFeatureImagesDir,
ensureAutomakerDir,
} from "../lib/automaker-paths.js";
import type { TokenUsage } from "../providers/types.js";
export interface Feature {
id: string;
@@ -43,6 +44,7 @@ export interface Feature {
error?: string;
summary?: string;
startedAt?: string;
tokenUsage?: TokenUsage;
[key: string]: unknown; // Keep catch-all for extensibility
}