fix: recompute token usage on reads

Author: doctarock

src/core/task/Task.ts

@@ -531,7 +531,6 @@ export class Task extends EventEmitter<TaskEvents> implements TaskLike {
 		this.usageTracker = new UsageTracker({
 			taskId: this.taskId,
 			getMessages: () => this.clineMessages,
-			getLastMessageTs: () => this.clineMessages.at(-1)?.ts,
 			emit: (event, ...args) => (this.emit as any)(event, ...args),
 		})
 
@@ -4649,7 +4648,7 @@ export class Task extends EventEmitter<TaskEvents> implements TaskLike {
 		return this.messageQueueService.messages
 	}
 
-	public get tokenUsage(): TokenUsage | undefined {
+	public get tokenUsage(): TokenUsage {
 		return this.usageTracker.tokenUsage
 	}
 

src/core/task/UsageTracker.ts

@@ -11,7 +11,6 @@ type EmitTaskEvent = (event: RooCodeEventName, ...args: any[]) => boolean
 export interface UsageTrackerOptions {
 	taskId: string
 	getMessages: () => ClineMessage[]
-	getLastMessageTs: () => number | undefined
 	emit: EmitTaskEvent
 	emitIntervalMs?: number
 }
@@ -21,25 +20,21 @@ const DEFAULT_TOKEN_USAGE_EMIT_INTERVAL_MS = 2000
 export class UsageTracker {
 	private readonly taskId: string
 	private readonly getMessages: () => ClineMessage[]
-	private readonly getLastMessageTs: () => number | undefined
 	private readonly emit: EmitTaskEvent
 	private readonly debouncedEmitTokenUsage: ReturnType<typeof debounce>
 
 	private tokenUsageSnapshot?: TokenUsage
-	private tokenUsageSnapshotAt?: number
 	private toolUsageSnapshot?: ToolUsage
 	private currentToolUsage: ToolUsage = {}
 
 	constructor({
 		taskId,
 		getMessages,
-		getLastMessageTs,
 		emit,
 		emitIntervalMs = DEFAULT_TOKEN_USAGE_EMIT_INTERVAL_MS,
 	}: UsageTrackerOptions) {
 		this.taskId = taskId
 		this.getMessages = getMessages
-		this.getLastMessageTs = getLastMessageTs
 		this.emit = emit
 
 		// Uses debounce with maxWait to achieve throttle-like behavior:
@@ -54,7 +49,6 @@ export class UsageTracker {
 				if (tokenChanged || toolChanged) {
 					this.emit(RooCodeEventName.TaskTokenUsageUpdated, this.taskId, tokenUsage, toolUsage)
 					this.tokenUsageSnapshot = tokenUsage
-					this.tokenUsageSnapshotAt = this.getLastMessageTs()
 					this.toolUsageSnapshot = JSON.parse(JSON.stringify(toolUsage))
 				}
 			},
@@ -102,15 +96,8 @@ export class UsageTracker {
 		this.currentToolUsage = toolUsage
 	}
 
-	public get tokenUsage(): TokenUsage | undefined {
-		if (this.tokenUsageSnapshot && this.tokenUsageSnapshotAt) {
-			return this.tokenUsageSnapshot
-		}
-
-		this.tokenUsageSnapshot = this.getTokenUsage()
-		this.tokenUsageSnapshotAt = this.getLastMessageTs()
-
-		return this.tokenUsageSnapshot
+	public get tokenUsage(): TokenUsage {
+		return this.getTokenUsage()
 	}
 
 	private ensureToolUsageEntry(toolName: ToolName): { attempts: number; failures: number } {

src/core/task/__tests__/UsageTracker.spec.ts

@@ -0,0 +1,56 @@
+import { describe, expect, it, vi } from "vitest"
+import { RooCodeEventName, type ClineMessage } from "@roo-code/types"
+
+import { UsageTracker } from "../UsageTracker.js"
+
+const createUsageMessage = (text: string, ts = 1000): ClineMessage => ({
+	type: "say",
+	say: "api_req_started",
+	text,
+	ts,
+})
+
+describe("UsageTracker", () => {
+	it("recomputes tokenUsage on each read after an emitted snapshot exists", () => {
+		const messages: ClineMessage[] = [
+			{ type: "say", say: "text", text: "task", ts: 1 },
+			createUsageMessage('{"tokensIn":100,"tokensOut":50,"cost":0.01}', 1000),
+		]
+		const tracker = new UsageTracker({
+			taskId: "task-id",
+			getMessages: () => messages,
+			emit: vi.fn(),
+		})
+
+		tracker.emitTokenUsageUpdate(tracker.getTokenUsage())
+		expect(tracker.tokenUsage.totalTokensIn).toBe(100)
+
+		messages[1] = createUsageMessage('{"tokensIn":250,"tokensOut":50,"cost":0.02}', 1000)
+
+		expect(tracker.tokenUsage.totalTokensIn).toBe(250)
+		expect(tracker.tokenUsage.totalCost).toBe(0.02)
+	})
+
+	it("keeps emitted snapshots for change comparisons", () => {
+		const emit = vi.fn()
+		const messages: ClineMessage[] = [
+			{ type: "say", say: "text", text: "task", ts: 1 },
+			createUsageMessage('{"tokensIn":100,"tokensOut":50,"cost":0.01}', 1000),
+		]
+		const tracker = new UsageTracker({
+			taskId: "task-id",
+			getMessages: () => messages,
+			emit,
+		})
+
+		tracker.emitTokenUsageUpdate(tracker.getTokenUsage())
+
+		expect(emit).toHaveBeenCalledWith(
+			RooCodeEventName.TaskTokenUsageUpdated,
+			"task-id",
+			expect.objectContaining({ totalTokensIn: 100 }),
+			{},
+		)
+		expect((tracker as any).tokenUsageSnapshot.totalTokensIn).toBe(100)
+	})
+})