fix(telemetry): narrow UseMcpToolTool callback, harden test mocks, close coverage gaps

Author: edelauna

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

@@ -15,14 +15,17 @@ vi.mock("../../../shared/modes", async (importOriginal) => {
 		getModeBySlug: vi.fn(actual.getModeBySlug),
 	}
 })
-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,
-		),
-	),
-}))
+// isValidToolName is left as the real implementation (only validateToolUse is
+// mocked): it has its own independent mcp_ prefix carve-out, and a hand-rolled
+// mock allowlist here would mask a regression in toTelemetryToolName's
+// ordering relative to isValidToolName.
+vi.mock("../../tools/validateToolUse", async (importOriginal) => {
+	const actual = await importOriginal<typeof import("../../tools/validateToolUse")>()
+	return {
+		...actual,
+		validateToolUse: vi.fn(),
+	}
+})
 
 vi.mock("@roo-code/core", () => ({
 	customToolRegistry: {

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

@@ -2,14 +2,18 @@
 
 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,
-		),
-	),
-}))
+// Only validateToolUse itself needs mocking (unused by toTelemetryToolName,
+// but imported by the same module). isValidToolName is left as the real
+// implementation: it has its own independent mcp_ prefix carve-out, and a
+// hand-rolled mock allowlist here would mask a regression where
+// toTelemetryToolName's ordering relative to isValidToolName changes.
+vi.mock("../../tools/validateToolUse", async (importOriginal) => {
+	const actual = await importOriginal<typeof import("../../tools/validateToolUse")>()
+	return {
+		...actual,
+		validateToolUse: vi.fn(),
+	}
+})
 
 import { toTelemetryToolName } from "../presentAssistantMessage"
 
@@ -40,4 +44,11 @@ describe("toTelemetryToolName", () => {
 		expect(result).not.toBe(raw)
 		expect(result).toBe("invalid_tool_call")
 	})
+
+	it("maps mcp_ names to use_mcp_tool against the real isValidToolName, not a mock allowlist", () => {
+		// isValidToolName is unmocked in this file (see the vi.mock factory above),
+		// so this exercises the real mcp_ prefix carve-out ordering rather than one
+		// a hand-rolled mock could silently keep agreeing with after a regression.
+		expect(toTelemetryToolName("mcp_my_server_do_thing", false, undefined)).toBe("use_mcp_tool")
+	})
 })

src/core/tools/BaseTool.ts

@@ -11,15 +11,6 @@ export interface ToolCallbacks {
 	handleError: HandleError
 	pushToolResult: PushToolResult
 	toolCallId?: string
-	/**
-	 * Optional hook invoked once a tool's own internal validation (params,
-	 * target existence, permissions) has passed, before side-effecting
-	 * execution begins. Used by callers that must defer telemetry attribution
-	 * until validation this deep can't be done from the outside (e.g. native
-	 * MCP tool calls, whose server/tool/allow-list checks live inside
-	 * UseMcpToolTool rather than the shared validateToolUse path).
-	 */
-	onValidated?: () => void
 }
 
 /**
@@ -117,9 +108,16 @@ export abstract class BaseTool<TName extends ToolName> {
 	 *
 	 * @param task - Task instance
 	 * @param block - ToolUse block from assistant message
-	 * @param callbacks - Tool execution callbacks
+	 * @param callbacks - Tool execution callbacks. Accepts any subclass-specific
+	 *   extension of ToolCallbacks (e.g. UseMcpToolCallbacks) so callers can pass
+	 *   extra fields through to a matching execute() override without widening
+	 *   the shared ToolCallbacks interface for every tool.
 	 */
-	async handle(task: Task, block: ToolUse<TName>, callbacks: ToolCallbacks): Promise<void> {
+	async handle<TCallbacks extends ToolCallbacks>(
+		task: Task,
+		block: ToolUse<TName>,
+		callbacks: TCallbacks,
+	): Promise<void> {
 		// Handle partial messages
 		if (block.partial) {
 			try {

src/core/tools/UseMcpToolTool.ts

@@ -15,6 +15,18 @@ interface UseMcpToolParams {
 	arguments?: Record<string, unknown>
 }
 
+/**
+ * Extends the shared callbacks with a hook invoked once this tool's own
+ * internal validation (params, tool existence, server allow-list) has
+ * passed, before side-effecting execution begins. Native MCP tool calls
+ * (presentAssistantMessage's mcp_tool_use branch) use this to defer telemetry
+ * attribution past validation that lives inside this class rather than the
+ * shared validateToolUse path.
+ */
+export interface UseMcpToolCallbacks extends ToolCallbacks {
+	onValidated?: () => void
+}
+
 type ValidationResult =
 	| { isValid: false }
 	| {
@@ -27,7 +39,7 @@ type ValidationResult =
 export class UseMcpToolTool extends BaseTool<"use_mcp_tool"> {
 	readonly name = "use_mcp_tool" as const
 
-	async execute(params: UseMcpToolParams, task: Task, callbacks: ToolCallbacks): Promise<void> {
+	async execute(params: UseMcpToolParams, task: Task, callbacks: UseMcpToolCallbacks): Promise<void> {
 		const { askApproval, handleError, pushToolResult, onValidated } = callbacks
 
 		try {

src/core/tools/__tests__/applyPatchTool.execute.spec.ts

@@ -0,0 +1,96 @@
+// npx vitest run core/tools/__tests__/applyPatchTool.execute.spec.ts
+
+import type { MockedFunction } from "vitest"
+
+import { fileExistsAtPath } from "../../../utils/fs"
+import { isPathOutsideWorkspace } from "../../../utils/pathUtils"
+import type { Task } from "../../task/Task"
+import { ApplyPatchTool } from "../ApplyPatchTool"
+
+vi.mock("fs/promises", () => ({
+	default: {
+		readFile: vi.fn().mockResolvedValue("original file content\n"),
+		unlink: vi.fn().mockResolvedValue(undefined),
+	},
+}))
+
+vi.mock("../../../utils/fs", () => ({
+	fileExistsAtPath: vi.fn().mockResolvedValue(true),
+}))
+
+vi.mock("../../../utils/pathUtils", () => ({
+	isPathOutsideWorkspace: vi.fn().mockReturnValue(false),
+}))
+
+describe("ApplyPatchTool.execute - delete file success path", () => {
+	const mockedFileExistsAtPath = fileExistsAtPath as MockedFunction<typeof fileExistsAtPath>
+	const mockedIsPathOutsideWorkspace = isPathOutsideWorkspace as MockedFunction<typeof isPathOutsideWorkspace>
+
+	let tool: ApplyPatchTool
+	let mockTask: Pick<
+		Task,
+		| "cwd"
+		| "consecutiveMistakeCount"
+		| "recordToolUsage"
+		| "recordToolError"
+		| "rooIgnoreController"
+		| "rooProtectedController"
+		| "say"
+		| "processQueuedMessages"
+		| "didEditFile"
+	>
+	let mockAskApproval: MockedFunction<(...args: unknown[]) => Promise<boolean>>
+	let mockHandleError: MockedFunction<(...args: unknown[]) => Promise<void>>
+	let mockPushToolResult: MockedFunction<(...args: unknown[]) => void>
+
+	beforeEach(() => {
+		vi.clearAllMocks()
+
+		mockedFileExistsAtPath.mockResolvedValue(true)
+		mockedIsPathOutsideWorkspace.mockReturnValue(false)
+
+		mockTask = {
+			cwd: "/workspace/project",
+			consecutiveMistakeCount: 0,
+			recordToolUsage: vi.fn(),
+			recordToolError: vi.fn(),
+			rooIgnoreController: {
+				validateAccess: vi.fn().mockReturnValue(true),
+			} as unknown as Task["rooIgnoreController"],
+			rooProtectedController: {
+				isWriteProtected: vi.fn().mockReturnValue(false),
+			} as unknown as Task["rooProtectedController"],
+			say: vi.fn().mockResolvedValue(undefined),
+			processQueuedMessages: vi.fn(),
+			didEditFile: false,
+		}
+
+		mockAskApproval = vi.fn().mockResolvedValue(true)
+		mockHandleError = vi.fn().mockResolvedValue(undefined)
+		mockPushToolResult = vi.fn()
+
+		tool = new ApplyPatchTool()
+	})
+
+	it("deletes the file and records no local tool usage on success", async () => {
+		const patch = `*** Begin Patch
+*** Delete File: src/obsolete.ts
+*** End Patch`
+
+		await tool.execute({ patch }, mockTask as Task, {
+			askApproval: mockAskApproval,
+			handleError: mockHandleError,
+			pushToolResult: mockPushToolResult,
+		})
+
+		expect(mockAskApproval).toHaveBeenCalled()
+		expect(mockPushToolResult).toHaveBeenCalledWith(expect.stringContaining("Successfully deleted"))
+		expect(mockTask.didEditFile).toBe(true)
+		expect(mockHandleError).not.toHaveBeenCalled()
+
+		// Usage is recorded once at the central presentAssistantMessage
+		// attribution point, not locally by the handler.
+		expect(mockTask.recordToolUsage).not.toHaveBeenCalled()
+		expect(mockTask.recordToolError).not.toHaveBeenCalled()
+	})
+})

src/core/tools/__tests__/generateImageTool.test.ts

@@ -163,6 +163,9 @@ describe("generateImageTool", () => {
 			expect(mockAskApproval).toHaveBeenCalled()
 			expect(mockGenerateImage).toHaveBeenCalled()
 			expect(mockPushToolResult).toHaveBeenCalled()
+			// Usage is recorded once at the central presentAssistantMessage
+			// attribution point, not locally by the handler.
+			expect(mockCline.recordToolUsage).not.toHaveBeenCalled()
 		})
 
 		it("should add cache-busting parameter to image URI", async () => {