Fix: Emit TaskTokenUsageUpdated when tool usage changes

hannesrudolph · hannesrudolph · commit f2ac0a11b08b · 2025-12-08T12:54:51.000-07:00
- Add hasToolUsageChanged() helper function to getApiMetrics.ts
- Add toolUsageSnapshot property to Task.ts to track tool usage changes
- Update saveClineMessages() to emit when either token OR tool usage changes
- Update emitFinalTokenUsageUpdate() to also check tool usage changes
- Add comprehensive tests for tool usage change detection

This ensures final tool usage stats are captured on task abort even if
token usage hasn't changed (e.g., when task is aborted before API request
completes but tools were already executed).
diff --git a/src/core/task/Task.ts b/src/core/task/Task.ts
@@ -61,7 +61,7 @@ import { combineApiRequests } from "../../shared/combineApiRequests"
 import { combineCommandSequences } from "../../shared/combineCommandSequences"
 import { t } from "../../i18n"
 import { ClineApiReqCancelReason, ClineApiReqInfo } from "../../shared/ExtensionMessage"
-import { getApiMetrics, hasTokenUsageChanged } from "../../shared/getApiMetrics"
+import { getApiMetrics, hasTokenUsageChanged, hasToolUsageChanged } from "../../shared/getApiMetrics"
 import { ClineAskResponse } from "../../shared/WebviewMessage"
 import { defaultModeSlug, getModeBySlug, getGroupName } from "../../shared/modes"
 import { DiffStrategy, type ToolUse, type ToolParamName, toolParamNames } from "../../shared/tools"
@@ -322,6 +322,9 @@ export class Task extends EventEmitter<TaskEvents> implements TaskLike {
 	private tokenUsageSnapshot?: TokenUsage
 	private tokenUsageSnapshotAt?: number
 
+	// Tool Usage Cache
+	private toolUsageSnapshot?: ToolUsage
+
 	// Token Usage Throttling
 	private lastTokenUsageEmitTime?: number
 	private readonly TOKEN_USAGE_EMIT_INTERVAL_MS = 2000 // 2 seconds
@@ -927,11 +930,17 @@ export class Task extends EventEmitter<TaskEvents> implements TaskLike {
 
 			const shouldEmitDueToThrottle = timeSinceLastEmit >= this.TOKEN_USAGE_EMIT_INTERVAL_MS
 
-			if (shouldEmitDueToThrottle && hasTokenUsageChanged(tokenUsage, this.tokenUsageSnapshot)) {
+			// Emit if throttle window allows AND either token usage or tool usage changed
+			const tokenChanged = hasTokenUsageChanged(tokenUsage, this.tokenUsageSnapshot)
+			const toolChanged = hasToolUsageChanged(this.toolUsage, this.toolUsageSnapshot)
+
+			if (shouldEmitDueToThrottle && (tokenChanged || toolChanged)) {
 				this.emit(RooCodeEventName.TaskTokenUsageUpdated, this.taskId, tokenUsage, this.toolUsage)
 				this.lastTokenUsageEmitTime = now
 				this.tokenUsageSnapshot = tokenUsage
 				this.tokenUsageSnapshotAt = this.clineMessages.at(-1)?.ts
+				// Deep copy tool usage for snapshot
+				this.toolUsageSnapshot = JSON.parse(JSON.stringify(this.toolUsage))
 			}
 
 			await this.providerRef.deref()?.updateTaskHistory(historyItem)
@@ -1856,11 +1865,16 @@ export class Task extends EventEmitter<TaskEvents> implements TaskLike {
 	 */
 	public emitFinalTokenUsageUpdate(): void {
 		const tokenUsage = this.getTokenUsage()
-		if (hasTokenUsageChanged(tokenUsage, this.tokenUsageSnapshot)) {
+		const tokenChanged = hasTokenUsageChanged(tokenUsage, this.tokenUsageSnapshot)
+		const toolChanged = hasToolUsageChanged(this.toolUsage, this.toolUsageSnapshot)
+
+		if (tokenChanged || toolChanged) {
 			this.emit(RooCodeEventName.TaskTokenUsageUpdated, this.taskId, tokenUsage, this.toolUsage)
 			this.tokenUsageSnapshot = tokenUsage
 			this.tokenUsageSnapshotAt = this.clineMessages.at(-1)?.ts
 			this.lastTokenUsageEmitTime = Date.now()
+			// Deep copy tool usage for snapshot
+			this.toolUsageSnapshot = JSON.parse(JSON.stringify(this.toolUsage))
 		}
 	}
 
diff --git a/src/core/task/__tests__/Task.throttle.test.ts b/src/core/task/__tests__/Task.throttle.test.ts
@@ -2,6 +2,7 @@ import { RooCodeEventName, ProviderSettings, TokenUsage, ToolUsage } from "@roo-
 
 import { Task } from "../Task"
 import { ClineProvider } from "../../webview/ClineProvider"
+import { hasToolUsageChanged, hasTokenUsageChanged } from "../../../shared/getApiMetrics"
 
 // Mock dependencies
 vi.mock("../../webview/ClineProvider")
@@ -381,4 +382,238 @@ describe("Task token usage throttling", () => {
 		// Should not have emitted again since token usage didn't change
 		expect(secondEmitCount).toBe(firstEmitCount)
 	})
+
+	test("should emit when tool usage changes even if token usage is the same", async () => {
+		const { taskMetadata } = await import("../../task-persistence")
+
+		// Mock taskMetadata to return same token usage
+		const constantTokenUsage: TokenUsage = {
+			totalTokensIn: 100,
+			totalTokensOut: 50,
+			totalCost: 0.01,
+			contextTokens: 150,
+			totalCacheWrites: 0,
+			totalCacheReads: 0,
+		}
+
+		vi.mocked(taskMetadata).mockResolvedValue({
+			historyItem: {
+				id: "test-task-id",
+				number: 1,
+				task: "Test task",
+				ts: Date.now(),
+				totalCost: 0.01,
+				tokensIn: 100,
+				tokensOut: 50,
+			},
+			tokenUsage: constantTokenUsage,
+		})
+
+		const emitSpy = vi.spyOn(task, "emit")
+
+		// Add first message - should emit
+		await (task as any).addToClineMessages({
+			ts: Date.now(),
+			type: "say",
+			say: "text",
+			text: "Message 1",
+		})
+
+		const firstEmitCount = emitSpy.mock.calls.filter(
+			(call) => call[0] === RooCodeEventName.TaskTokenUsageUpdated,
+		).length
+
+		// Wait for throttle period
+		vi.advanceTimersByTime(2100)
+
+		// Change tool usage (token usage stays the same)
+		task.toolUsage = {
+			read_file: { attempts: 5, failures: 1 },
+		}
+
+		// Add another message
+		await (task as any).addToClineMessages({
+			ts: Date.now(),
+			type: "say",
+			say: "text",
+			text: "Message 2",
+		})
+
+		const secondEmitCount = emitSpy.mock.calls.filter(
+			(call) => call[0] === RooCodeEventName.TaskTokenUsageUpdated,
+		).length
+
+		// Should have emitted because tool usage changed even though token usage didn't
+		expect(secondEmitCount).toBeGreaterThan(firstEmitCount)
+	})
+
+	test("should update toolUsageSnapshot when emission occurs", async () => {
+		// Add initial message
+		await (task as any).addToClineMessages({
+			ts: Date.now(),
+			type: "say",
+			say: "text",
+			text: "Message 1",
+		})
+
+		// Initially toolUsageSnapshot should be set to current toolUsage (empty object)
+		const initialSnapshot = (task as any).toolUsageSnapshot
+		expect(initialSnapshot).toBeDefined()
+		expect(Object.keys(initialSnapshot)).toHaveLength(0)
+
+		// Wait for throttle period
+		vi.advanceTimersByTime(2100)
+
+		// Update tool usage
+		task.toolUsage = {
+			read_file: { attempts: 3, failures: 0 },
+			write_to_file: { attempts: 2, failures: 1 },
+		}
+
+		// Add another message
+		await (task as any).addToClineMessages({
+			ts: Date.now(),
+			type: "say",
+			say: "text",
+			text: "Message 2",
+		})
+
+		// Snapshot should be updated to match the new toolUsage
+		const newSnapshot = (task as any).toolUsageSnapshot
+		expect(newSnapshot).not.toBe(initialSnapshot)
+		expect(newSnapshot.read_file).toEqual({ attempts: 3, failures: 0 })
+		expect(newSnapshot.write_to_file).toEqual({ attempts: 2, failures: 1 })
+	})
+
+	test("emitFinalTokenUsageUpdate should emit on tool usage change alone", async () => {
+		const emitSpy = vi.spyOn(task, "emit")
+
+		// Set initial tool usage and simulate previous emission
+		;(task as any).tokenUsageSnapshot = task.getTokenUsage()
+		;(task as any).toolUsageSnapshot = {}
+
+		// Change tool usage
+		task.toolUsage = {
+			execute_command: { attempts: 1, failures: 0 },
+		}
+
+		// Call emitFinalTokenUsageUpdate
+		task.emitFinalTokenUsageUpdate()
+
+		// Should emit due to tool usage change
+		expect(emitSpy).toHaveBeenCalledWith(
+			RooCodeEventName.TaskTokenUsageUpdated,
+			task.taskId,
+			expect.any(Object),
+			task.toolUsage,
+		)
+	})
+})
+
+describe("hasToolUsageChanged", () => {
+	test("should return true when snapshot is undefined and current has data", () => {
+		const current: ToolUsage = {
+			read_file: { attempts: 1, failures: 0 },
+		}
+		expect(hasToolUsageChanged(current, undefined)).toBe(true)
+	})
+
+	test("should return false when both are empty", () => {
+		expect(hasToolUsageChanged({}, {})).toBe(false)
+	})
+
+	test("should return false when snapshot is undefined and current is empty", () => {
+		expect(hasToolUsageChanged({}, undefined)).toBe(false)
+	})
+
+	test("should return true when a new tool is added", () => {
+		const current: ToolUsage = {
+			read_file: { attempts: 1, failures: 0 },
+			write_to_file: { attempts: 1, failures: 0 },
+		}
+		const snapshot: ToolUsage = {
+			read_file: { attempts: 1, failures: 0 },
+		}
+		expect(hasToolUsageChanged(current, snapshot)).toBe(true)
+	})
+
+	test("should return true when attempts change", () => {
+		const current: ToolUsage = {
+			read_file: { attempts: 2, failures: 0 },
+		}
+		const snapshot: ToolUsage = {
+			read_file: { attempts: 1, failures: 0 },
+		}
+		expect(hasToolUsageChanged(current, snapshot)).toBe(true)
+	})
+
+	test("should return true when failures change", () => {
+		const current: ToolUsage = {
+			read_file: { attempts: 1, failures: 1 },
+		}
+		const snapshot: ToolUsage = {
+			read_file: { attempts: 1, failures: 0 },
+		}
+		expect(hasToolUsageChanged(current, snapshot)).toBe(true)
+	})
+
+	test("should return false when nothing changed", () => {
+		const current: ToolUsage = {
+			read_file: { attempts: 3, failures: 1 },
+			write_to_file: { attempts: 2, failures: 0 },
+		}
+		const snapshot: ToolUsage = {
+			read_file: { attempts: 3, failures: 1 },
+			write_to_file: { attempts: 2, failures: 0 },
+		}
+		expect(hasToolUsageChanged(current, snapshot)).toBe(false)
+	})
+})
+
+describe("hasTokenUsageChanged", () => {
+	test("should return true when snapshot is undefined", () => {
+		const current: TokenUsage = {
+			totalTokensIn: 100,
+			totalTokensOut: 50,
+			totalCost: 0.01,
+			contextTokens: 150,
+		}
+		expect(hasTokenUsageChanged(current, undefined)).toBe(true)
+	})
+
+	test("should return true when totalTokensIn changes", () => {
+		const current: TokenUsage = {
+			totalTokensIn: 200,
+			totalTokensOut: 50,
+			totalCost: 0.01,
+			contextTokens: 150,
+		}
+		const snapshot: TokenUsage = {
+			totalTokensIn: 100,
+			totalTokensOut: 50,
+			totalCost: 0.01,
+			contextTokens: 150,
+		}
+		expect(hasTokenUsageChanged(current, snapshot)).toBe(true)
+	})
+
+	test("should return false when nothing changed", () => {
+		const current: TokenUsage = {
+			totalTokensIn: 100,
+			totalTokensOut: 50,
+			totalCost: 0.01,
+			contextTokens: 150,
+			totalCacheWrites: 10,
+			totalCacheReads: 5,
+		}
+		const snapshot: TokenUsage = {
+			totalTokensIn: 100,
+			totalTokensOut: 50,
+			totalCost: 0.01,
+			contextTokens: 150,
+			totalCacheWrites: 10,
+			totalCacheReads: 5,
+		}
+		expect(hasTokenUsageChanged(current, snapshot)).toBe(false)
+	})
 })
diff --git a/src/shared/getApiMetrics.ts b/src/shared/getApiMetrics.ts
@@ -1,4 +1,4 @@
-import type { TokenUsage, ClineMessage } from "@roo-code/types"
+import type { TokenUsage, ToolUsage, ToolName, ClineMessage } from "@roo-code/types"
 
 export type ParsedApiReqStartedTextType = {
 	tokensIn: number
@@ -123,3 +123,35 @@ export function hasTokenUsageChanged(current: TokenUsage, snapshot?: TokenUsage)
 
 	return keysToCompare.some((key) => current[key] !== snapshot[key])
 }
+
+/**
+ * Check if tool usage has changed by comparing attempts and failures.
+ * @param current - Current tool usage data
+ * @param snapshot - Previous snapshot to compare against
+ * @returns true if any tool's attempts/failures have changed or snapshot is undefined
+ */
+export function hasToolUsageChanged(current: ToolUsage, snapshot?: ToolUsage): boolean {
+	if (!snapshot) {
+		return Object.keys(current).length > 0
+	}
+
+	const currentKeys = Object.keys(current) as ToolName[]
+	const snapshotKeys = Object.keys(snapshot) as ToolName[]
+
+	// Check if number of tools changed
+	if (currentKeys.length !== snapshotKeys.length) {
+		return true
+	}
+
+	// Check if any tool's stats changed
+	return currentKeys.some((key) => {
+		const currentTool = current[key]
+		const snapshotTool = snapshot[key]
+
+		if (!snapshotTool || !currentTool) {
+			return true
+		}
+
+		return currentTool.attempts !== snapshotTool.attempts || currentTool.failures !== snapshotTool.failures
+	})
+}
