fix(telemetry): defer native MCP usage recording until validation passes

Author: edelauna

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

@@ -4,9 +4,17 @@ 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 { getModeBySlug } from "../../../shared/modes"
 import type { Task } from "../../task/Task"
 
 vi.mock("../../task/Task")
+vi.mock("../../../shared/modes", async (importOriginal) => {
+	const actual = await importOriginal<typeof import("../../../shared/modes")>()
+	return {
+		...actual,
+		getModeBySlug: vi.fn(actual.getModeBySlug),
+	}
+})
 vi.mock("../../tools/validateToolUse", () => ({
 	validateToolUse: vi.fn(),
 	isValidToolName: vi.fn((toolName: string) =>
@@ -53,7 +61,12 @@ interface MockTask {
 	recordToolUsage: ReturnType<typeof vi.fn>
 	recordToolError: ReturnType<typeof vi.fn>
 	toolRepetitionDetector: { check: ReturnType<typeof vi.fn> }
-	providerRef: { deref: () => { getState: ReturnType<typeof vi.fn> } }
+	providerRef: {
+		deref: () => {
+			getState: ReturnType<typeof vi.fn>
+			getMcpHub?: () => { findServerNameBySanitizedName: (name: string) => string | undefined }
+		}
+	}
 	say: ReturnType<typeof vi.fn>
 	ask: ReturnType<typeof vi.fn>
 	pushToolResultToUserContent: ReturnType<typeof vi.fn>
@@ -218,4 +231,80 @@ describe("presentAssistantMessage - tool usage attribution", () => {
 		expect(mockTask.recordToolError).not.toHaveBeenCalledWith("totally_made_up_tool", expect.anything())
 		expect(mockTask.recordToolUsage).not.toHaveBeenCalled()
 	})
+
+	describe("native mcp_tool_use block", () => {
+		it("records exactly one attempt once the MCP tool's own validation passes", async () => {
+			mockTask.providerRef = {
+				deref: () => ({
+					getState: vi.fn().mockResolvedValue({
+						mode: "code",
+						customModes: [],
+					}),
+					getMcpHub: () => ({
+						findServerNameBySanitizedName: () => "my_server",
+					}),
+				}),
+			}
+
+			mockTask.assistantMessageContent = [
+				{
+					type: "mcp_tool_use",
+					id: "call_native_mcp",
+					name: "mcp_my_server_do_thing",
+					serverName: "my_server",
+					toolName: "do_thing",
+					arguments: {},
+					partial: false,
+				},
+			]
+
+			await presentAssistantMessage(mockTask as unknown as Task)
+
+			expect(mockTask.recordToolUsage).toHaveBeenCalledTimes(1)
+			expect(mockTask.recordToolUsage).toHaveBeenCalledWith("use_mcp_tool")
+			expect(TelemetryService.instance.captureToolUsage).toHaveBeenCalledTimes(1)
+			expect(TelemetryService.instance.captureToolUsage).toHaveBeenCalledWith(mockTask.taskId, "use_mcp_tool")
+		})
+
+		it("records no attempt when the MCP server is not on the mode's allow-list", async () => {
+			vi.mocked(getModeBySlug).mockReturnValueOnce({
+				slug: "code",
+				name: "Code",
+				roleDefinition: "",
+				groups: [],
+				allowedMcpServers: ["some-other-server"],
+			})
+
+			mockTask.providerRef = {
+				deref: () => ({
+					getState: vi.fn().mockResolvedValue({
+						mode: "code",
+						customModes: [],
+					}),
+					getMcpHub: () => ({
+						findServerNameBySanitizedName: () => "my_server",
+					}),
+				}),
+			}
+
+			mockTask.assistantMessageContent = [
+				{
+					type: "mcp_tool_use",
+					id: "call_native_mcp_disallowed",
+					name: "mcp_my_server_do_thing",
+					serverName: "my_server",
+					toolName: "do_thing",
+					arguments: {},
+					partial: false,
+				},
+			]
+
+			await presentAssistantMessage(mockTask as unknown as Task)
+
+			// The server is disallowed, so the call never reaches onValidated:
+			// no success attempt is recorded for a call that was never permitted to execute.
+			expect(mockTask.recordToolUsage).not.toHaveBeenCalled()
+			expect(TelemetryService.instance.captureToolUsage).not.toHaveBeenCalled()
+		})
+	})
 })

src/core/assistant-message/presentAssistantMessage.ts

@@ -258,11 +258,6 @@ export async function presentAssistantMessage(cline: Task) {
 				pushToolResult(formatResponse.toolError(errorString))
 			}
 
-			if (!mcpBlock.partial) {
-				cline.recordToolUsage("use_mcp_tool") // Record as use_mcp_tool for analytics
-				TelemetryService.instance.captureToolUsage(cline.taskId, "use_mcp_tool")
-			}
-
 			// Resolve sanitized server name back to original server name
 			// The serverName from parsing is sanitized (e.g., "my_server" from "my server")
 			// We need the original name to find the actual MCP connection
@@ -298,6 +293,12 @@ export async function presentAssistantMessage(cline: Task) {
 				askApproval,
 				handleError,
 				pushToolResult,
+				onValidated: mcpBlock.partial
+					? undefined
+					: () => {
+							cline.recordToolUsage("use_mcp_tool")
+							TelemetryService.instance.captureToolUsage(cline.taskId, "use_mcp_tool")
+						},
 			})
 			break
 		}

src/core/tools/BaseTool.ts

@@ -11,6 +11,15 @@ 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
 }
 
 /**

src/core/tools/UseMcpToolTool.ts

@@ -28,7 +28,7 @@ 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> {
-		const { askApproval, handleError, pushToolResult } = callbacks
+		const { askApproval, handleError, pushToolResult, onValidated } = callbacks
 
 		try {
 			// Validate parameters
@@ -66,6 +66,10 @@ export class UseMcpToolTool extends BaseTool<"use_mcp_tool"> {
 			// Reset mistake count on successful validation
 			task.consecutiveMistakeCount = 0
 
+			// All internal validation (params, tool existence, server allow-list) has
+			// passed. Only now is it safe to attribute this as an attempted tool use.
+			onValidated?.()
+
 			// Get user approval
 			const completeMessage = JSON.stringify({
 				type: "use_mcp_tool",