Skip to content

Commit c138f3d

Browse files
committed
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
1 parent 72668fe commit c138f3d

File tree

10 files changed

+508
-1
lines changed

10 files changed

+508
-1
lines changed

packages/types/src/index.ts

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@ export * from "./telemetry.js"
2020
export * from "./terminal.js"
2121
export * from "./tool.js"
2222
export * from "./type-fu.js"
23+
export * from "./usage.js"
2324
export * from "./vscode.js"
2425

2526
export * from "./providers/index.js"

packages/types/src/usage.ts

Lines changed: 42 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,42 @@
1+
import { z } from "zod"
2+
3+
/**
4+
* Schema for persistent usage tracking item
5+
* Stores cost and token data for completed tasks
6+
*/
7+
export const PersistentUsageItemSchema = z.object({
8+
taskId: z.string(),
9+
taskWorkspace: z.string(), // Workspace path for filtering
10+
mode: z.string().optional(),
11+
timestamp: z.number(),
12+
cost: z.number().default(0),
13+
inputTokens: z.number().default(0),
14+
outputTokens: z.number().default(0),
15+
cacheReads: z.number().default(0),
16+
cacheWrites: z.number().default(0),
17+
})
18+
19+
export type PersistentUsageItem = z.infer<typeof PersistentUsageItemSchema>
20+
21+
/**
22+
* Usage summary data structure for frontend display
23+
*/
24+
export interface UsageSummary {
25+
totalCost: number
26+
totalInputTokens: number
27+
totalOutputTokens: number
28+
totalCacheReads: number
29+
totalCacheWrites: number
30+
modeBreakdown: Record<
31+
string,
32+
{
33+
cost: number
34+
inputTokens: number
35+
outputTokens: number
36+
cacheReads: number
37+
cacheWrites: number
38+
count: number
39+
}
40+
>
41+
workspaceName?: string
42+
}

src/core/webview/ClineProvider.ts

Lines changed: 46 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1432,6 +1432,42 @@ export class ClineProvider
14321432
if (!this.checkMdmCompliance()) {
14331433
await this.postMessageToWebview({ type: "action", action: "accountButtonClicked" })
14341434
}
1435+
1436+
// Load persistent usage data asynchronously
1437+
this.loadPersistentUsageDataAsync()
1438+
}
1439+
1440+
/**
1441+
* Load persistent usage data asynchronously and send to webview
1442+
*/
1443+
public async loadPersistentUsageDataAsync() {
1444+
try {
1445+
const { getUsageTrackingService } = await import("../../services/usage-tracking")
1446+
const usageService = getUsageTrackingService()
1447+
1448+
// Get both all workspaces and current workspace data
1449+
const [allData, currentData] = await Promise.all([
1450+
usageService.getUsageSummary(),
1451+
usageService.getUsageSummary(this.cwd),
1452+
])
1453+
1454+
// Send to webview
1455+
await this.postMessageToWebview({
1456+
type: "persistentUsageData",
1457+
persistentUsageData: {
1458+
all: allData,
1459+
current: currentData,
1460+
},
1461+
})
1462+
1463+
// Migrate existing task history if needed
1464+
const taskHistory = this.getGlobalState("taskHistory") ?? []
1465+
if (taskHistory.length > 0) {
1466+
await usageService.migrateFromTaskHistory(taskHistory, this.cwd)
1467+
}
1468+
} catch (error) {
1469+
this.log(`Failed to load persistent usage data: ${error}`)
1470+
}
14351471
}
14361472

14371473
/**
@@ -1957,6 +1993,16 @@ export class ClineProvider
19571993
}
19581994

19591995
await this.updateGlobalState("taskHistory", history)
1996+
1997+
// Update persistent usage tracking
1998+
try {
1999+
const { getUsageTrackingService } = await import("../../services/usage-tracking")
2000+
const usageService = getUsageTrackingService()
2001+
await usageService.updateUsageTracking(item, this.cwd)
2002+
} catch (error) {
2003+
this.log(`Failed to update usage tracking: ${error}`)
2004+
}
2005+
19602006
return history
19612007
}
19622008

src/core/webview/webviewMessageHandler.ts

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2604,5 +2604,10 @@ export const webviewMessageHandler = async (
26042604
}
26052605
break
26062606
}
2607+
case "requestPersistentUsageData": {
2608+
// Load persistent usage data asynchronously
2609+
provider.loadPersistentUsageDataAsync()
2610+
break
2611+
}
26072612
}
26082613
}
Lines changed: 250 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,250 @@
1+
import * as path from "path"
2+
import * as os from "os"
3+
import { PersistentUsageItem, PersistentUsageItemSchema, UsageSummary, HistoryItem } from "@roo-code/types"
4+
import { safeWriteJson } from "../../utils/safeWriteJson"
5+
import { fileExistsAtPath } from "../../utils/fs"
6+
import * as fs from "fs/promises"
7+
8+
/**
9+
* Service for managing persistent usage tracking data
10+
* Stores data in ~/.roo/usage-tracking.json
11+
*/
12+
export class UsageTrackingService {
13+
private static instance: UsageTrackingService | null = null
14+
private readonly filePath: string
15+
16+
private constructor() {
17+
// Store in global ~/.roo directory
18+
const rooDir = path.join(os.homedir(), ".roo")
19+
this.filePath = path.join(rooDir, "usage-tracking.json")
20+
}
21+
22+
public static getInstance(): UsageTrackingService {
23+
if (!UsageTrackingService.instance) {
24+
UsageTrackingService.instance = new UsageTrackingService()
25+
}
26+
return UsageTrackingService.instance
27+
}
28+
29+
/**
30+
* Load existing usage data from disk
31+
*/
32+
private async loadData(): Promise<PersistentUsageItem[]> {
33+
try {
34+
if (!(await fileExistsAtPath(this.filePath))) {
35+
return []
36+
}
37+
38+
const content = await fs.readFile(this.filePath, "utf-8")
39+
const data = JSON.parse(content)
40+
41+
// Validate and parse each item
42+
if (Array.isArray(data)) {
43+
return data
44+
.map((item) => {
45+
try {
46+
return PersistentUsageItemSchema.parse(item)
47+
} catch {
48+
// Skip invalid items
49+
return null
50+
}
51+
})
52+
.filter((item): item is PersistentUsageItem => item !== null)
53+
}
54+
55+
return []
56+
} catch (error) {
57+
console.error("Failed to load usage tracking data:", error)
58+
return []
59+
}
60+
}
61+
62+
/**
63+
* Save usage data to disk
64+
*/
65+
private async saveData(data: PersistentUsageItem[]): Promise<void> {
66+
try {
67+
await safeWriteJson(this.filePath, data)
68+
} catch (error) {
69+
console.error("Failed to save usage tracking data:", error)
70+
throw error
71+
}
72+
}
73+
74+
/**
75+
* Update usage tracking with new task data
76+
*/
77+
public async updateUsageTracking(historyItem: HistoryItem, workspacePath: string): Promise<void> {
78+
try {
79+
const data = await this.loadData()
80+
81+
// Remove existing entry for this task if it exists
82+
const filteredData = data.filter((item) => item.taskId !== historyItem.id)
83+
84+
// Calculate total cost and tokens from the history item
85+
let totalCost = 0
86+
let totalInputTokens = 0
87+
let totalOutputTokens = 0
88+
let totalCacheReads = 0
89+
let totalCacheWrites = 0
90+
91+
// Sum up all API request info from the task
92+
if (historyItem.totalCost) {
93+
totalCost = historyItem.totalCost
94+
}
95+
if (historyItem.tokensIn) {
96+
totalInputTokens = historyItem.tokensIn
97+
}
98+
if (historyItem.tokensOut) {
99+
totalOutputTokens = historyItem.tokensOut
100+
}
101+
if (historyItem.cacheReads) {
102+
totalCacheReads = historyItem.cacheReads
103+
}
104+
if (historyItem.cacheWrites) {
105+
totalCacheWrites = historyItem.cacheWrites
106+
}
107+
108+
// Add new usage item
109+
const newItem: PersistentUsageItem = {
110+
taskId: historyItem.id,
111+
taskWorkspace: workspacePath,
112+
mode: historyItem.mode,
113+
timestamp: historyItem.ts,
114+
cost: totalCost,
115+
inputTokens: totalInputTokens,
116+
outputTokens: totalOutputTokens,
117+
cacheReads: totalCacheReads,
118+
cacheWrites: totalCacheWrites,
119+
}
120+
121+
filteredData.push(newItem)
122+
123+
// Save updated data
124+
await this.saveData(filteredData)
125+
} catch (error) {
126+
console.error("Failed to update usage tracking:", error)
127+
// Don't throw - we don't want to break the main flow
128+
}
129+
}
130+
131+
/**
132+
* Get usage summary for display
133+
*/
134+
public async getUsageSummary(workspacePath?: string): Promise<UsageSummary> {
135+
try {
136+
const data = await this.loadData()
137+
138+
// Filter by workspace if specified
139+
const filteredData = workspacePath ? data.filter((item) => item.taskWorkspace === workspacePath) : data
140+
141+
// Calculate totals
142+
let totalCost = 0
143+
let totalInputTokens = 0
144+
let totalOutputTokens = 0
145+
let totalCacheReads = 0
146+
let totalCacheWrites = 0
147+
const modeBreakdown: UsageSummary["modeBreakdown"] = {}
148+
149+
for (const item of filteredData) {
150+
totalCost += item.cost
151+
totalInputTokens += item.inputTokens
152+
totalOutputTokens += item.outputTokens
153+
totalCacheReads += item.cacheReads
154+
totalCacheWrites += item.cacheWrites
155+
156+
// Update mode breakdown
157+
const mode = item.mode || "unknown"
158+
if (!modeBreakdown[mode]) {
159+
modeBreakdown[mode] = {
160+
cost: 0,
161+
inputTokens: 0,
162+
outputTokens: 0,
163+
cacheReads: 0,
164+
cacheWrites: 0,
165+
count: 0,
166+
}
167+
}
168+
169+
modeBreakdown[mode].cost += item.cost
170+
modeBreakdown[mode].inputTokens += item.inputTokens
171+
modeBreakdown[mode].outputTokens += item.outputTokens
172+
modeBreakdown[mode].cacheReads += item.cacheReads
173+
modeBreakdown[mode].cacheWrites += item.cacheWrites
174+
modeBreakdown[mode].count += 1
175+
}
176+
177+
// Get workspace name from path
178+
const workspaceName = workspacePath ? path.basename(workspacePath) : undefined
179+
180+
return {
181+
totalCost,
182+
totalInputTokens,
183+
totalOutputTokens,
184+
totalCacheReads,
185+
totalCacheWrites,
186+
modeBreakdown,
187+
workspaceName,
188+
}
189+
} catch (error) {
190+
console.error("Failed to get usage summary:", error)
191+
// Return empty summary on error
192+
return {
193+
totalCost: 0,
194+
totalInputTokens: 0,
195+
totalOutputTokens: 0,
196+
totalCacheReads: 0,
197+
totalCacheWrites: 0,
198+
modeBreakdown: {},
199+
}
200+
}
201+
}
202+
203+
/**
204+
* Migrate existing task history to persistent storage
205+
*/
206+
public async migrateFromTaskHistory(taskHistory: HistoryItem[], workspacePath: string): Promise<void> {
207+
try {
208+
const data = await this.loadData()
209+
const existingTaskIds = new Set(data.map((item) => item.taskId))
210+
211+
// Only migrate tasks that aren't already in persistent storage
212+
const tasksToMigrate = taskHistory.filter((task) => !existingTaskIds.has(task.id))
213+
214+
for (const task of tasksToMigrate) {
215+
if (task.totalCost || task.tokensIn || task.tokensOut) {
216+
await this.updateUsageTracking(task, workspacePath)
217+
}
218+
}
219+
} catch (error) {
220+
console.error("Failed to migrate task history:", error)
221+
}
222+
}
223+
224+
/**
225+
* Clear all usage data (for testing or reset)
226+
*/
227+
public async clearAllData(): Promise<void> {
228+
try {
229+
await this.saveData([])
230+
} catch (error) {
231+
console.error("Failed to clear usage data:", error)
232+
}
233+
}
234+
235+
/**
236+
* Remove usage data for a specific task
237+
*/
238+
public async removeTask(taskId: string): Promise<void> {
239+
try {
240+
const data = await this.loadData()
241+
const filteredData = data.filter((item) => item.taskId !== taskId)
242+
await this.saveData(filteredData)
243+
} catch (error) {
244+
console.error("Failed to remove task from usage tracking:", error)
245+
}
246+
}
247+
}
248+
249+
// Export singleton instance getter
250+
export const getUsageTrackingService = () => UsageTrackingService.getInstance()

src/shared/ExtensionMessage.ts

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -121,6 +121,7 @@ export interface ExtensionMessage {
121121
| "showEditMessageDialog"
122122
| "commands"
123123
| "insertTextIntoTextarea"
124+
| "persistentUsageData"
124125
text?: string
125126
payload?: any // Add a generic payload for now, can refine later
126127
action?:
@@ -195,6 +196,10 @@ export interface ExtensionMessage {
195196
messageTs?: number
196197
context?: string
197198
commands?: Command[]
199+
persistentUsageData?: {
200+
all: any // UsageSummary type
201+
current: any // UsageSummary type
202+
}
198203
}
199204

200205
export type ExtensionState = Pick<

src/shared/WebviewMessage.ts

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -210,6 +210,7 @@ export interface WebviewMessage {
210210
| "deleteCommand"
211211
| "createCommand"
212212
| "insertTextIntoTextarea"
213+
| "requestPersistentUsageData"
213214
text?: string
214215
editedMessageContent?: string
215216
tab?: "settings" | "history" | "mcp" | "modes" | "chat" | "marketplace" | "account"

0 commit comments

Comments
 (0)