From c138f3dfecfa0dbc3500f4288d5610ce72064095 Mon Sep 17 00:00:00 2001 From: Roo Code Date: Thu, 7 Aug 2025 19:34:12 +0000 Subject: [PATCH] feat: add AI Inference Summary for total costs and token usage - Add persistent usage tracking service that stores data in ~/.roo/usage-tracking.json - Create CostSummary React component to display costs and token breakdown - Integrate usage tracking with task history updates - Add workspace filtering (Current vs All Workspaces) - Implement mode breakdown showing costs per AI mode - Data persists even when tasks are deleted from history Fixes #6822 --- packages/types/src/index.ts | 1 + packages/types/src/usage.ts | 42 +++ src/core/webview/ClineProvider.ts | 46 ++++ src/core/webview/webviewMessageHandler.ts | 5 + src/services/usage-tracking/index.ts | 250 ++++++++++++++++++ src/shared/ExtensionMessage.ts | 5 + src/shared/WebviewMessage.ts | 1 + .../src/components/welcome/CostSummary.tsx | 133 ++++++++++ .../src/components/welcome/WelcomeView.tsx | 8 +- .../src/context/ExtensionStateContext.tsx | 18 ++ 10 files changed, 508 insertions(+), 1 deletion(-) create mode 100644 packages/types/src/usage.ts create mode 100644 src/services/usage-tracking/index.ts create mode 100644 webview-ui/src/components/welcome/CostSummary.tsx diff --git a/packages/types/src/index.ts b/packages/types/src/index.ts index dcbb1c4f54..0fe36d7a89 100644 --- a/packages/types/src/index.ts +++ b/packages/types/src/index.ts @@ -20,6 +20,7 @@ export * from "./telemetry.js" export * from "./terminal.js" export * from "./tool.js" export * from "./type-fu.js" +export * from "./usage.js" export * from "./vscode.js" export * from "./providers/index.js" diff --git a/packages/types/src/usage.ts b/packages/types/src/usage.ts new file mode 100644 index 0000000000..db34ddab56 --- /dev/null +++ b/packages/types/src/usage.ts @@ -0,0 +1,42 @@ +import { z } from "zod" + +/** + * Schema for persistent usage tracking item + * Stores cost and token data for completed tasks + */ +export const PersistentUsageItemSchema = z.object({ + taskId: z.string(), + taskWorkspace: z.string(), // Workspace path for filtering + mode: z.string().optional(), + timestamp: z.number(), + cost: z.number().default(0), + inputTokens: z.number().default(0), + outputTokens: z.number().default(0), + cacheReads: z.number().default(0), + cacheWrites: z.number().default(0), +}) + +export type PersistentUsageItem = z.infer + +/** + * Usage summary data structure for frontend display + */ +export interface UsageSummary { + totalCost: number + totalInputTokens: number + totalOutputTokens: number + totalCacheReads: number + totalCacheWrites: number + modeBreakdown: Record< + string, + { + cost: number + inputTokens: number + outputTokens: number + cacheReads: number + cacheWrites: number + count: number + } + > + workspaceName?: string +} diff --git a/src/core/webview/ClineProvider.ts b/src/core/webview/ClineProvider.ts index 274060a19b..cb11017c86 100644 --- a/src/core/webview/ClineProvider.ts +++ b/src/core/webview/ClineProvider.ts @@ -1432,6 +1432,42 @@ export class ClineProvider if (!this.checkMdmCompliance()) { await this.postMessageToWebview({ type: "action", action: "accountButtonClicked" }) } + + // Load persistent usage data asynchronously + this.loadPersistentUsageDataAsync() + } + + /** + * Load persistent usage data asynchronously and send to webview + */ + public async loadPersistentUsageDataAsync() { + try { + const { getUsageTrackingService } = await import("../../services/usage-tracking") + const usageService = getUsageTrackingService() + + // Get both all workspaces and current workspace data + const [allData, currentData] = await Promise.all([ + usageService.getUsageSummary(), + usageService.getUsageSummary(this.cwd), + ]) + + // Send to webview + await this.postMessageToWebview({ + type: "persistentUsageData", + persistentUsageData: { + all: allData, + current: currentData, + }, + }) + + // Migrate existing task history if needed + const taskHistory = this.getGlobalState("taskHistory") ?? [] + if (taskHistory.length > 0) { + await usageService.migrateFromTaskHistory(taskHistory, this.cwd) + } + } catch (error) { + this.log(`Failed to load persistent usage data: ${error}`) + } } /** @@ -1957,6 +1993,16 @@ export class ClineProvider } await this.updateGlobalState("taskHistory", history) + + // Update persistent usage tracking + try { + const { getUsageTrackingService } = await import("../../services/usage-tracking") + const usageService = getUsageTrackingService() + await usageService.updateUsageTracking(item, this.cwd) + } catch (error) { + this.log(`Failed to update usage tracking: ${error}`) + } + return history } diff --git a/src/core/webview/webviewMessageHandler.ts b/src/core/webview/webviewMessageHandler.ts index e2c6d6a475..219463ca1d 100644 --- a/src/core/webview/webviewMessageHandler.ts +++ b/src/core/webview/webviewMessageHandler.ts @@ -2604,5 +2604,10 @@ export const webviewMessageHandler = async ( } break } + case "requestPersistentUsageData": { + // Load persistent usage data asynchronously + provider.loadPersistentUsageDataAsync() + break + } } } diff --git a/src/services/usage-tracking/index.ts b/src/services/usage-tracking/index.ts new file mode 100644 index 0000000000..3d7223b749 --- /dev/null +++ b/src/services/usage-tracking/index.ts @@ -0,0 +1,250 @@ +import * as path from "path" +import * as os from "os" +import { PersistentUsageItem, PersistentUsageItemSchema, UsageSummary, HistoryItem } from "@roo-code/types" +import { safeWriteJson } from "../../utils/safeWriteJson" +import { fileExistsAtPath } from "../../utils/fs" +import * as fs from "fs/promises" + +/** + * Service for managing persistent usage tracking data + * Stores data in ~/.roo/usage-tracking.json + */ +export class UsageTrackingService { + private static instance: UsageTrackingService | null = null + private readonly filePath: string + + private constructor() { + // Store in global ~/.roo directory + const rooDir = path.join(os.homedir(), ".roo") + this.filePath = path.join(rooDir, "usage-tracking.json") + } + + public static getInstance(): UsageTrackingService { + if (!UsageTrackingService.instance) { + UsageTrackingService.instance = new UsageTrackingService() + } + return UsageTrackingService.instance + } + + /** + * Load existing usage data from disk + */ + private async loadData(): Promise { + try { + if (!(await fileExistsAtPath(this.filePath))) { + return [] + } + + const content = await fs.readFile(this.filePath, "utf-8") + const data = JSON.parse(content) + + // Validate and parse each item + if (Array.isArray(data)) { + return data + .map((item) => { + try { + return PersistentUsageItemSchema.parse(item) + } catch { + // Skip invalid items + return null + } + }) + .filter((item): item is PersistentUsageItem => item !== null) + } + + return [] + } catch (error) { + console.error("Failed to load usage tracking data:", error) + return [] + } + } + + /** + * Save usage data to disk + */ + private async saveData(data: PersistentUsageItem[]): Promise { + try { + await safeWriteJson(this.filePath, data) + } catch (error) { + console.error("Failed to save usage tracking data:", error) + throw error + } + } + + /** + * Update usage tracking with new task data + */ + public async updateUsageTracking(historyItem: HistoryItem, workspacePath: string): Promise { + try { + const data = await this.loadData() + + // Remove existing entry for this task if it exists + const filteredData = data.filter((item) => item.taskId !== historyItem.id) + + // Calculate total cost and tokens from the history item + let totalCost = 0 + let totalInputTokens = 0 + let totalOutputTokens = 0 + let totalCacheReads = 0 + let totalCacheWrites = 0 + + // Sum up all API request info from the task + if (historyItem.totalCost) { + totalCost = historyItem.totalCost + } + if (historyItem.tokensIn) { + totalInputTokens = historyItem.tokensIn + } + if (historyItem.tokensOut) { + totalOutputTokens = historyItem.tokensOut + } + if (historyItem.cacheReads) { + totalCacheReads = historyItem.cacheReads + } + if (historyItem.cacheWrites) { + totalCacheWrites = historyItem.cacheWrites + } + + // Add new usage item + const newItem: PersistentUsageItem = { + taskId: historyItem.id, + taskWorkspace: workspacePath, + mode: historyItem.mode, + timestamp: historyItem.ts, + cost: totalCost, + inputTokens: totalInputTokens, + outputTokens: totalOutputTokens, + cacheReads: totalCacheReads, + cacheWrites: totalCacheWrites, + } + + filteredData.push(newItem) + + // Save updated data + await this.saveData(filteredData) + } catch (error) { + console.error("Failed to update usage tracking:", error) + // Don't throw - we don't want to break the main flow + } + } + + /** + * Get usage summary for display + */ + public async getUsageSummary(workspacePath?: string): Promise { + try { + const data = await this.loadData() + + // Filter by workspace if specified + const filteredData = workspacePath ? data.filter((item) => item.taskWorkspace === workspacePath) : data + + // Calculate totals + let totalCost = 0 + let totalInputTokens = 0 + let totalOutputTokens = 0 + let totalCacheReads = 0 + let totalCacheWrites = 0 + const modeBreakdown: UsageSummary["modeBreakdown"] = {} + + for (const item of filteredData) { + totalCost += item.cost + totalInputTokens += item.inputTokens + totalOutputTokens += item.outputTokens + totalCacheReads += item.cacheReads + totalCacheWrites += item.cacheWrites + + // Update mode breakdown + const mode = item.mode || "unknown" + if (!modeBreakdown[mode]) { + modeBreakdown[mode] = { + cost: 0, + inputTokens: 0, + outputTokens: 0, + cacheReads: 0, + cacheWrites: 0, + count: 0, + } + } + + modeBreakdown[mode].cost += item.cost + modeBreakdown[mode].inputTokens += item.inputTokens + modeBreakdown[mode].outputTokens += item.outputTokens + modeBreakdown[mode].cacheReads += item.cacheReads + modeBreakdown[mode].cacheWrites += item.cacheWrites + modeBreakdown[mode].count += 1 + } + + // Get workspace name from path + const workspaceName = workspacePath ? path.basename(workspacePath) : undefined + + return { + totalCost, + totalInputTokens, + totalOutputTokens, + totalCacheReads, + totalCacheWrites, + modeBreakdown, + workspaceName, + } + } catch (error) { + console.error("Failed to get usage summary:", error) + // Return empty summary on error + return { + totalCost: 0, + totalInputTokens: 0, + totalOutputTokens: 0, + totalCacheReads: 0, + totalCacheWrites: 0, + modeBreakdown: {}, + } + } + } + + /** + * Migrate existing task history to persistent storage + */ + public async migrateFromTaskHistory(taskHistory: HistoryItem[], workspacePath: string): Promise { + try { + const data = await this.loadData() + const existingTaskIds = new Set(data.map((item) => item.taskId)) + + // Only migrate tasks that aren't already in persistent storage + const tasksToMigrate = taskHistory.filter((task) => !existingTaskIds.has(task.id)) + + for (const task of tasksToMigrate) { + if (task.totalCost || task.tokensIn || task.tokensOut) { + await this.updateUsageTracking(task, workspacePath) + } + } + } catch (error) { + console.error("Failed to migrate task history:", error) + } + } + + /** + * Clear all usage data (for testing or reset) + */ + public async clearAllData(): Promise { + try { + await this.saveData([]) + } catch (error) { + console.error("Failed to clear usage data:", error) + } + } + + /** + * Remove usage data for a specific task + */ + public async removeTask(taskId: string): Promise { + try { + const data = await this.loadData() + const filteredData = data.filter((item) => item.taskId !== taskId) + await this.saveData(filteredData) + } catch (error) { + console.error("Failed to remove task from usage tracking:", error) + } + } +} + +// Export singleton instance getter +export const getUsageTrackingService = () => UsageTrackingService.getInstance() diff --git a/src/shared/ExtensionMessage.ts b/src/shared/ExtensionMessage.ts index 3ddd69945c..7df2232021 100644 --- a/src/shared/ExtensionMessage.ts +++ b/src/shared/ExtensionMessage.ts @@ -121,6 +121,7 @@ export interface ExtensionMessage { | "showEditMessageDialog" | "commands" | "insertTextIntoTextarea" + | "persistentUsageData" text?: string payload?: any // Add a generic payload for now, can refine later action?: @@ -195,6 +196,10 @@ export interface ExtensionMessage { messageTs?: number context?: string commands?: Command[] + persistentUsageData?: { + all: any // UsageSummary type + current: any // UsageSummary type + } } export type ExtensionState = Pick< diff --git a/src/shared/WebviewMessage.ts b/src/shared/WebviewMessage.ts index cb8759d851..a451cb683e 100644 --- a/src/shared/WebviewMessage.ts +++ b/src/shared/WebviewMessage.ts @@ -210,6 +210,7 @@ export interface WebviewMessage { | "deleteCommand" | "createCommand" | "insertTextIntoTextarea" + | "requestPersistentUsageData" text?: string editedMessageContent?: string tab?: "settings" | "history" | "mcp" | "modes" | "chat" | "marketplace" | "account" diff --git a/webview-ui/src/components/welcome/CostSummary.tsx b/webview-ui/src/components/welcome/CostSummary.tsx new file mode 100644 index 0000000000..3e38a4e76d --- /dev/null +++ b/webview-ui/src/components/welcome/CostSummary.tsx @@ -0,0 +1,133 @@ +import React, { useState, useEffect } from "react" +import { VSCodeButton } from "@vscode/webview-ui-toolkit/react" +import { UsageSummary } from "@roo-code/types" +import { vscode } from "@src/utils/vscode" + +interface CostSummaryProps { + usageData?: { + all: UsageSummary + current: UsageSummary + } +} + +const CostSummary: React.FC = ({ usageData }) => { + const [viewMode, setViewMode] = useState<"all" | "current">("all") + + // Request usage data on mount + useEffect(() => { + vscode.postMessage({ type: "requestPersistentUsageData" }) + }, []) + + if (!usageData) { + return ( +
+

AI Inference Summary

+
Loading usage data...
+
+ ) + } + + const currentData = viewMode === "all" ? usageData.all : usageData.current + const workspaceName = viewMode === "current" ? usageData.current.workspaceName : "All Workspaces" + + // Format currency + const formatCost = (cost: number) => { + return `$${cost.toFixed(4)}` + } + + // Format large numbers with commas + const formatNumber = (num: number) => { + return num.toLocaleString() + } + + return ( +
+
+

AI Inference Summary

+
+ setViewMode("current")} + className="text-xs"> + Current + + setViewMode("all")} + className="text-xs"> + All + +
+
+ +
{workspaceName}
+ + {/* Total Summary */} +
+
+
Total Cost
+
{formatCost(currentData.totalCost)}
+
+
+
Total Tokens
+
+ {formatNumber(currentData.totalInputTokens + currentData.totalOutputTokens)} +
+
+
+ + {/* Token Breakdown */} +
+
+ Input: + {formatNumber(currentData.totalInputTokens)} +
+
+ Output: + {formatNumber(currentData.totalOutputTokens)} +
+ {(currentData.totalCacheReads > 0 || currentData.totalCacheWrites > 0) && ( + <> +
+ Cache Reads: + {formatNumber(currentData.totalCacheReads)} +
+
+ Cache Writes: + {formatNumber(currentData.totalCacheWrites)} +
+ + )} +
+ + {/* Mode Breakdown */} + {Object.keys(currentData.modeBreakdown).length > 0 && ( +
+
Breakdown by Mode
+
+ {Object.entries(currentData.modeBreakdown) + .sort(([, a], [, b]) => b.cost - a.cost) + .map(([mode, data]) => ( +
+
+ {mode} + + ({data.count} task{data.count !== 1 ? "s" : ""}) + +
+
+ {formatCost(data.cost)} + + {formatNumber(data.inputTokens + data.outputTokens)} tokens + +
+
+ ))} +
+
+ )} +
+ ) +} + +export default CostSummary diff --git a/webview-ui/src/components/welcome/WelcomeView.tsx b/webview-ui/src/components/welcome/WelcomeView.tsx index 215a290d0d..7886a7c9c0 100644 --- a/webview-ui/src/components/welcome/WelcomeView.tsx +++ b/webview-ui/src/components/welcome/WelcomeView.tsx @@ -15,9 +15,11 @@ import ApiOptions from "../settings/ApiOptions" import { Tab, TabContent } from "../common/Tab" import RooHero from "./RooHero" +import CostSummary from "./CostSummary" const WelcomeView = () => { - const { apiConfiguration, currentApiConfigName, setApiConfiguration, uriScheme, machineId } = useExtensionState() + const { apiConfiguration, currentApiConfigName, setApiConfiguration, uriScheme, machineId, persistentUsageData } = + useExtensionState() const { t } = useAppTranslation() const [errorMessage, setErrorMessage] = useState(undefined) @@ -51,6 +53,10 @@ const WelcomeView = () => { + + {/* Add CostSummary component */} + +

{t("welcome:greeting")}

diff --git a/webview-ui/src/context/ExtensionStateContext.tsx b/webview-ui/src/context/ExtensionStateContext.tsx index da7ab63358..a2dbf0282d 100644 --- a/webview-ui/src/context/ExtensionStateContext.tsx +++ b/webview-ui/src/context/ExtensionStateContext.tsx @@ -147,6 +147,10 @@ export interface ExtensionStateContextType extends ExtensionState { setMaxDiagnosticMessages: (value: number) => void includeTaskHistoryInEnhance?: boolean setIncludeTaskHistoryInEnhance: (value: boolean) => void + persistentUsageData?: { + all: any // UsageSummary type + current: any // UsageSummary type + } } export const ExtensionStateContext = createContext(undefined) @@ -266,6 +270,13 @@ export const ExtensionStateContextProvider: React.FC<{ children: React.ReactNode global: {}, }) const [includeTaskHistoryInEnhance, setIncludeTaskHistoryInEnhance] = useState(false) + const [persistentUsageData, setPersistentUsageData] = useState< + | { + all: any // UsageSummary type + current: any // UsageSummary type + } + | undefined + >(undefined) const setListApiConfigMeta = useCallback( (value: ProviderSettingsEntry[]) => setState((prevState) => ({ ...prevState, listApiConfigMeta: value })), @@ -369,6 +380,12 @@ export const ExtensionStateContextProvider: React.FC<{ children: React.ReactNode } break } + case "persistentUsageData": { + if (message.persistentUsageData) { + setPersistentUsageData(message.persistentUsageData) + } + break + } } }, [setListApiConfigMeta], @@ -395,6 +412,7 @@ export const ExtensionStateContextProvider: React.FC<{ children: React.ReactNode filePaths, openedTabs, commands, + persistentUsageData, soundVolume: state.soundVolume, ttsSpeed: state.ttsSpeed, fuzzyMatchThreshold: state.fuzzyMatchThreshold,