fix(telemetry): record tool usage once centrally, sanitize raw tool names

Author: edelauna

packages/types/src/tool.ts

@@ -46,6 +46,7 @@ export const toolNames = [
 	"skill",
 	"generate_image",
 	"custom_tool",
+	"invalid_tool_call",
 ] as const
 
 export const toolNamesSchema = z.enum(toolNames)

src/core/assistant-message/__tests__/presentAssistantMessage-tool-usage-attribution.spec.ts

@@ -0,0 +1,221 @@
+// npx vitest src/core/assistant-message/__tests__/presentAssistantMessage-tool-usage-attribution.spec.ts
+
+import type { Anthropic } from "@anthropic-ai/sdk"
+import { describe, it, expect, beforeEach, vi } from "vitest"
+import { presentAssistantMessage } from "../presentAssistantMessage"
+import { validateToolUse } from "../../tools/validateToolUse"
+import type { Task } from "../../task/Task"
+
+vi.mock("../../task/Task")
+vi.mock("../../tools/validateToolUse", () => ({
+	validateToolUse: vi.fn(),
+	isValidToolName: vi.fn((toolName: string) =>
+		["read_file", "write_to_file", "ask_followup_question", "attempt_completion", "use_mcp_tool"].includes(
+			toolName,
+		),
+	),
+}))
+
+vi.mock("@roo-code/core", () => ({
+	customToolRegistry: {
+		has: vi.fn(() => false),
+		get: vi.fn(),
+	},
+}))
+
+vi.mock("@roo-code/telemetry", () => ({
+	TelemetryService: {
+		instance: {
+			captureToolUsage: vi.fn(),
+			captureConsecutiveMistakeError: vi.fn(),
+			captureEvent: vi.fn(),
+		},
+	},
+}))
+
+import { TelemetryService } from "@roo-code/telemetry"
+
+interface MockTask {
+	taskId: string
+	instanceId: string
+	abort: boolean
+	presentAssistantMessageLocked: boolean
+	presentAssistantMessageHasPendingUpdates: boolean
+	currentStreamingContentIndex: number
+	assistantMessageContent: unknown[]
+	userMessageContent: Anthropic.ToolResultBlockParam[]
+	didCompleteReadingStream: boolean
+	didRejectTool: boolean
+	didAlreadyUseTool: boolean
+	consecutiveMistakeCount: number
+	clineMessages: unknown[]
+	api: { getModel: () => { id: string; info: Record<string, unknown> } }
+	recordToolUsage: ReturnType<typeof vi.fn>
+	recordToolError: ReturnType<typeof vi.fn>
+	toolRepetitionDetector: { check: ReturnType<typeof vi.fn> }
+	providerRef: { deref: () => { getState: ReturnType<typeof vi.fn> } }
+	say: ReturnType<typeof vi.fn>
+	ask: ReturnType<typeof vi.fn>
+	pushToolResultToUserContent: ReturnType<typeof vi.fn>
+}
+
+describe("presentAssistantMessage - tool usage attribution", () => {
+	let mockTask: MockTask
+
+	beforeEach(() => {
+		vi.clearAllMocks()
+		vi.mocked(validateToolUse).mockImplementation(() => undefined)
+
+		mockTask = {
+			taskId: "test-task-id",
+			instanceId: "test-instance",
+			abort: false,
+			presentAssistantMessageLocked: false,
+			presentAssistantMessageHasPendingUpdates: false,
+			currentStreamingContentIndex: 0,
+			assistantMessageContent: [],
+			userMessageContent: [],
+			didCompleteReadingStream: false,
+			didRejectTool: false,
+			didAlreadyUseTool: false,
+			consecutiveMistakeCount: 0,
+			clineMessages: [],
+			api: {
+				getModel: () => ({ id: "test-model", info: {} }),
+			},
+			recordToolUsage: vi.fn(),
+			recordToolError: vi.fn(),
+			toolRepetitionDetector: {
+				check: vi.fn().mockReturnValue({ allowExecution: true }),
+			},
+			providerRef: {
+				deref: () => ({
+					getState: vi.fn().mockResolvedValue({
+						mode: "code",
+						customModes: [],
+					}),
+				}),
+			},
+			say: vi.fn().mockResolvedValue(undefined),
+			ask: vi.fn().mockResolvedValue({ response: "yesButtonClicked" }),
+			pushToolResultToUserContent: vi.fn(),
+		}
+
+		mockTask.pushToolResultToUserContent = vi
+			.fn()
+			.mockImplementation((toolResult: Anthropic.ToolResultBlockParam) => {
+				const existingResult = mockTask.userMessageContent.find(
+					(block) => block.type === "tool_result" && block.tool_use_id === toolResult.tool_use_id,
+				)
+				if (existingResult) {
+					return false
+				}
+				mockTask.userMessageContent.push(toolResult)
+				return true
+			})
+	})
+
+	it("records exactly one attempt for a normal static tool", async () => {
+		mockTask.assistantMessageContent = [
+			{
+				type: "tool_use",
+				id: "call_1",
+				name: "read_file",
+				params: { path: "test.txt" },
+				nativeArgs: { path: "test.txt" },
+				partial: false,
+			},
+		]
+
+		await presentAssistantMessage(mockTask as unknown as Task)
+
+		expect(mockTask.recordToolUsage).toHaveBeenCalledTimes(1)
+		expect(mockTask.recordToolUsage).toHaveBeenCalledWith("read_file")
+		expect(TelemetryService.instance.captureToolUsage).toHaveBeenCalledTimes(1)
+		expect(TelemetryService.instance.captureToolUsage).toHaveBeenCalledWith(mockTask.taskId, "read_file")
+	})
+
+	it("records a valid dynamic mcp_ tool name as use_mcp_tool", async () => {
+		mockTask.assistantMessageContent = [
+			{
+				type: "tool_use",
+				id: "call_mcp",
+				name: "mcp_my_server_do_thing",
+				params: {},
+				nativeArgs: {},
+				partial: false,
+			},
+		]
+
+		await presentAssistantMessage(mockTask as unknown as Task)
+
+		expect(mockTask.recordToolUsage).toHaveBeenCalledWith("use_mcp_tool")
+		expect(TelemetryService.instance.captureToolUsage).toHaveBeenCalledWith(mockTask.taskId, "use_mcp_tool")
+	})
+
+	it("records a malformed mcp_ tool name as use_mcp_tool, not the raw name", async () => {
+		mockTask.assistantMessageContent = [
+			{
+				type: "tool_use",
+				id: "call_mcp_bad",
+				name: "mcp_",
+				params: {},
+				nativeArgs: {},
+				partial: false,
+			},
+		]
+
+		await presentAssistantMessage(mockTask as unknown as Task)
+
+		expect(mockTask.recordToolUsage).toHaveBeenCalledWith("use_mcp_tool")
+		expect(mockTask.recordToolUsage).not.toHaveBeenCalledWith("mcp_")
+	})
+
+	it("records a safe failure key without leaking the raw tool name when validation fails", async () => {
+		vi.mocked(validateToolUse).mockImplementation(() => {
+			throw new Error('Tool "read_file" is not allowed in this mode.')
+		})
+
+		mockTask.assistantMessageContent = [
+			{
+				type: "tool_use",
+				id: "call_bad_mode",
+				name: "read_file",
+				params: { path: "test.txt" },
+				nativeArgs: { path: "test.txt" },
+				partial: false,
+			},
+		]
+
+		await presentAssistantMessage(mockTask as unknown as Task)
+
+		// A known static tool that fails validation still maps to its own name
+		// (it's a real, recognized tool - just disallowed here), never left raw/unmapped.
+		expect(mockTask.recordToolError).toHaveBeenCalledWith("read_file", expect.any(String))
+		// No success attempt should be recorded for a validation failure.
+		expect(mockTask.recordToolUsage).not.toHaveBeenCalled()
+	})
+
+	it("records invalid_tool_call, not the raw name, when an arbitrary unknown tool fails validation", async () => {
+		vi.mocked(validateToolUse).mockImplementation(() => {
+			throw new Error('Unknown tool "totally_made_up_tool". This tool does not exist.')
+		})
+
+		mockTask.assistantMessageContent = [
+			{
+				type: "tool_use",
+				id: "call_unknown",
+				name: "totally_made_up_tool",
+				params: {},
+				nativeArgs: {},
+				partial: false,
+			},
+		]
+
+		await presentAssistantMessage(mockTask as unknown as Task)
+
+		expect(mockTask.recordToolError).toHaveBeenCalledWith("invalid_tool_call", expect.any(String))
+		expect(mockTask.recordToolError).not.toHaveBeenCalledWith("totally_made_up_tool", expect.anything())
+		expect(mockTask.recordToolUsage).not.toHaveBeenCalled()
+	})
+})

src/core/assistant-message/__tests__/presentAssistantMessage-unknown-tool.spec.ts

@@ -101,9 +101,10 @@ describe("presentAssistantMessage - Unknown Tool Handling", () => {
 		// Verify consecutiveMistakeCount was incremented
 		expect(mockTask.consecutiveMistakeCount).toBe(1)
 
-		// Verify recordToolError was called
+		// Verify recordToolError was called with a safe static key, never the
+		// raw model-controlled tool name.
 		expect(mockTask.recordToolError).toHaveBeenCalledWith(
-			"nonexistent_tool",
+			"invalid_tool_call",
 			expect.stringContaining("Unknown tool"),
 		)
 
@@ -135,8 +136,9 @@ describe("presentAssistantMessage - Unknown Tool Handling", () => {
 		// Verify consecutiveMistakeCount was incremented
 		expect(mockTask.consecutiveMistakeCount).toBe(1)
 
-		// Verify recordToolError was called
-		expect(mockTask.recordToolError).toHaveBeenCalled()
+		// Verify recordToolError was called with a safe static key, never the
+		// raw model-reported tool name ("fake_tool_that_does_not_exist").
+		expect(mockTask.recordToolError).toHaveBeenCalledWith("invalid_tool_call", expect.anything())
 
 		// Verify error message was shown to user
 		expect(mockTask.say).toHaveBeenCalledWith("error", expect.anything())

src/core/assistant-message/__tests__/toTelemetryToolName.spec.ts

@@ -0,0 +1,43 @@
+// npx vitest src/core/assistant-message/__tests__/toTelemetryToolName.spec.ts
+
+import { describe, it, expect, vi } from "vitest"
+
+vi.mock("../../tools/validateToolUse", () => ({
+	validateToolUse: vi.fn(),
+	isValidToolName: vi.fn((toolName: string) =>
+		["read_file", "write_to_file", "ask_followup_question", "attempt_completion", "use_mcp_tool"].includes(
+			toolName,
+		),
+	),
+}))
+
+import { toTelemetryToolName } from "../presentAssistantMessage"
+
+describe("toTelemetryToolName", () => {
+	it("maps a known static tool to its own name", () => {
+		expect(toTelemetryToolName("read_file", false, undefined)).toBe("read_file")
+	})
+
+	it("maps a registered custom tool to custom_tool", () => {
+		expect(toTelemetryToolName("my_custom_tool", true, undefined)).toBe("custom_tool")
+	})
+
+	it("maps a valid dynamic mcp_ tool name to use_mcp_tool", () => {
+		expect(toTelemetryToolName("mcp_my_server_do_thing", false, undefined)).toBe("use_mcp_tool")
+	})
+
+	it("maps a malformed mcp_ tool name to use_mcp_tool", () => {
+		expect(toTelemetryToolName("mcp_", false, undefined)).toBe("use_mcp_tool")
+	})
+
+	it("maps an arbitrary unknown tool name to invalid_tool_call", () => {
+		expect(toTelemetryToolName("drop_table_users", false, undefined)).toBe("invalid_tool_call")
+	})
+
+	it("never returns the raw name for an unrecognized tool", () => {
+		const raw = "'; DROP TABLE users; --"
+		const result = toTelemetryToolName(raw, false, undefined)
+		expect(result).not.toBe(raw)
+		expect(result).toBe("invalid_tool_call")
+	})
+})