ros/automaker
1
0
Fork 0
mirror of https://github.com/AutoMaker-Org/automaker.git synced 2026-01-30 06:12:03 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat: add token consumption tracking per task

Browse Source 
- 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
Download Patch File Download Diff File Expand all files Collapse all files

32
apps/server/src/providers/types.ts
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>;
}
/**

56
apps/server/src/services/auto-mode-service.ts
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(

2
apps/server/src/services/feature-loader.ts
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
}

62
apps/ui/src/components/ui/log-viewer.tsx
View File

@@ -24,6 +24,7 @@ import {
Circle,
Play,
Loader2,
Coins,
} from "lucide-react";
import { cn } from "@/lib/utils";
import {
@@ -34,10 +35,35 @@ import {
type LogEntryType,
type ToolCategory,
} from "@/lib/log-parser";
import type { TokenUsage } from "@/store/app-store";
interface LogViewerProps {
output: string;
className?: string;
tokenUsage?: TokenUsage;
}
/**
* Formats token counts for compact display (e.g., 12500 -> "12.5K")
*/
function formatTokenCount(count: number): string {
if (count >= 1000000) {
return `${(count / 1000000).toFixed(1)}M`;
}
if (count >= 1000) {
return `${(count / 1000).toFixed(1)}K`;
}
return count.toString();
}
/**
* Formats cost for display (e.g., 0.0847 -> "$0.0847")
*/
function formatCost(cost: number): string {
if (cost < 0.01) {
return `$${cost.toFixed(4)}`;
}
return `$${cost.toFixed(2)}`;
}
const getLogIcon = (type: LogEntryType) => {
@@ -413,7 +439,7 @@ interface ToolCategoryStats {
other: number;
}
export function LogViewer({ output, className }: LogViewerProps) {
export function LogViewer({ output, className, tokenUsage }: LogViewerProps) {
const [expandedIds, setExpandedIds] = useState<Set<string>>(new Set());
const [searchQuery, setSearchQuery] = useState("");
const [hiddenTypes, setHiddenTypes] = useState<Set<LogEntryType>>(new Set());
@@ -615,6 +641,40 @@ export function LogViewer({ output, className }: LogViewerProps) {
return (
<div className={cn("flex flex-col", className)}>
{/* Token Usage Summary Header */}
{tokenUsage && tokenUsage.totalTokens > 0 && (
<div className="mb-3 p-2 bg-zinc-900/50 rounded-lg border border-zinc-700/50">
<div className="flex items-center gap-3 text-xs text-muted-foreground">
<span className="flex items-center gap-1.5">
<Coins className="w-3.5 h-3.5 text-amber-400" />
<span className="font-medium">{formatTokenCount(tokenUsage.totalTokens)}</span>
<span className="text-muted-foreground/60">tokens</span>
</span>
<span className="text-muted-foreground/30">|</span>
<span className="flex items-center gap-1">
<span className="text-green-400">IN:</span>
<span>{formatTokenCount(tokenUsage.inputTokens)}</span>
</span>
<span className="flex items-center gap-1">
<span className="text-blue-400">OUT:</span>
<span>{formatTokenCount(tokenUsage.outputTokens)}</span>
</span>
{tokenUsage.cacheReadInputTokens > 0 && (
<>
<span className="text-muted-foreground/30">|</span>
<span className="flex items-center gap-1">
<span className="text-purple-400">Cache:</span>
<span>{formatTokenCount(tokenUsage.cacheReadInputTokens)}</span>
</span>
</>
)}
<span className="text-muted-foreground/30">|</span>
<span className="flex items-center gap-1 text-amber-400 font-medium">
{formatCost(tokenUsage.costUSD)}
</span>
</div>
</div>
)}
{/* Sticky header with search, stats, and filters */}
{/* Use -top-4 to compensate for parent's p-4 padding, pt-4 to restore visual spacing */}
<div className="sticky -top-4 z-10 bg-zinc-950/95 backdrop-blur-sm pt-4 pb-2 space-y-2 -mx-4 px-4">

54
apps/ui/src/components/views/board-view/components/kanban-card.tsx
View File

@@ -57,6 +57,7 @@ import {
Wand2,
Archive,
Lock,
Coins,
} from "lucide-react";
import { CountUpTimer } from "@/components/ui/count-up-timer";
import { getElectronAPI } from "@/lib/electron";
@@ -90,6 +91,29 @@ function formatThinkingLevel(level: ThinkingLevel | undefined): string {
return labels[level];
}
/**
* Formats token counts for compact display (e.g., 12500 -> "12.5K")
*/
function formatTokenCount(count: number): string {
if (count >= 1000000) {
return `${(count / 1000000).toFixed(1)}M`;
}
if (count >= 1000) {
return `${(count / 1000).toFixed(1)}K`;
}
return count.toString();
}
/**
* Formats cost for display (e.g., 0.0847 -> "$0.0847")
*/
function formatCost(cost: number): string {
if (cost < 0.01) {
return `$${cost.toFixed(4)}`;
}
return `$${cost.toFixed(2)}`;
}
interface KanbanCardProps {
feature: Feature;
onEdit: () => void;
@@ -873,6 +897,36 @@ export const KanbanCard = memo(function KanbanCard({
)}
</div>
)}
{/* Token Usage Display */}
{feature.tokenUsage && feature.tokenUsage.totalTokens > 0 && (
<div className="flex items-center gap-2 text-[10px] text-muted-foreground/60 pt-2 border-t border-border/30">
<TooltipProvider delayDuration={200}>
<Tooltip>
<TooltipTrigger asChild>
<span className="flex items-center gap-1 cursor-help">
<Coins className="w-2.5 h-2.5 text-amber-400" />
{formatTokenCount(feature.tokenUsage.totalTokens)} tokens
<span className="text-amber-400/80">
({formatCost(feature.tokenUsage.costUSD)})
</span>
</span>
</TooltipTrigger>
<TooltipContent side="top" className="text-xs">
<div className="space-y-1">
<p>Input: {formatTokenCount(feature.tokenUsage.inputTokens)}</p>
<p>Output: {formatTokenCount(feature.tokenUsage.outputTokens)}</p>
{feature.tokenUsage.cacheReadInputTokens > 0 && (
<p>Cache read: {formatTokenCount(feature.tokenUsage.cacheReadInputTokens)}</p>
)}
<p className="font-medium pt-1 border-t border-border/30">
Cost: {formatCost(feature.tokenUsage.costUSD)}
</p>
</div>
</TooltipContent>
</Tooltip>
</TooltipProvider>
</div>
)}
</>
)}
</div>

4
apps/ui/src/components/views/board-view/dialogs/agent-output-modal.tsx
View File

@@ -44,6 +44,8 @@ export function AgentOutputModal({
const autoScrollRef = useRef(true);
const projectPathRef = useRef<string>("");
const useWorktrees = useAppStore((state) => state.useWorktrees);
const features = useAppStore((state) => state.features);
const feature = features.find((f) => f.id === featureId);
// Auto-scroll to bottom when output changes
useEffect(() => {
@@ -387,7 +389,7 @@ export function AgentOutputModal({
No output yet. The agent will stream output here as it works.
</div>
) : viewMode === "parsed" ? (
<LogViewer output={output} />
<LogViewer output={output} tokenUsage={feature?.tokenUsage} />
) : (
<div className="whitespace-pre-wrap break-words text-zinc-300">
{output}

11
apps/ui/src/store/app-store.ts
View File

@@ -278,6 +278,16 @@ export interface AIProfile {
icon?: string; // Optional icon name from lucide
}
// Token usage statistics for tracking API costs per task
export interface TokenUsage {
inputTokens: number;
outputTokens: number;
cacheReadInputTokens: number;
cacheCreationInputTokens: number;
totalTokens: number;
costUSD: number;
}
export interface Feature {
id: string;
category: string;
@@ -305,6 +315,7 @@ export interface Feature {
planningMode?: PlanningMode; // Planning mode for this feature
planSpec?: PlanSpec; // Generated spec/plan data
requirePlanApproval?: boolean; // Whether to pause and require manual approval before implementation
tokenUsage?: TokenUsage; // Token consumption tracking for this task
}
// Parsed task from spec (for spec and full planning modes)