mirror of
https://github.com/AutoMaker-Org/automaker.git
synced 2026-02-01 08:13:37 +00:00
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:
@@ -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>;
|
||||
}
|
||||
|
||||
/**
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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
|
||||
}
|
||||
|
||||
|
||||
Reference in New Issue
Block a user